Base reverse mode for allocating functions on split (#149)

* Base reverse mode for allocating functions on split

* No fallback structure

Author: gdalle

DifferentiationInterface/docs/src/design.md

@@ -24,29 +24,5 @@ For simplicity, we remove `value_` in the operator names below.
     Full edges in the following graphs require a single call to the destination.
     Dotted edges require multiple calls to the destination, the number is indicated on the edge.
 
-### First order
-
-```mermaid
-flowchart LR
-    direction LR
-    subgraph Out-of-place
-        pushforward
-        pullback
-        derivative --> pushforward
-        gradient --> pullback
-        jacobian .-> |n|pushforward
-        jacobian .-> |m|pullback
-    end
-
-    subgraph In-place
-        pushforward!! --> pushforward
-        pullback!! --> pullback
-        derivative!! --> pushforward!!
-        gradient!! --> pullback!!
-        jacobian!! .-> |n|pushforward!!
-        jacobian!! .-> |m|pullback!!
-    end
-
-    pushforward .-> |m|pullback
-    pullback .-> |n|pushforward
-```
+!!! warning
+    This is still in flux, come back later!

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -19,13 +19,6 @@ DI.mode(::AutoReverseChainRules) = ADTypes.AbstractReverseMode
 
 DI.prepare_pullback(f, ::AutoForwardChainRules, x) = NoPullbackExtras()
 
-function DI.value_and_pullback(f, backend::AutoReverseChainRules, x, dy, ::NoPullbackExtras)
-    rc = ruleconfig(backend)
-    y, pullback = rrule_via_ad(rc, f, x)
-    _, new_dx = pullback(dy)
-    return y, new_dx
-end
-
 function DI.value_and_pullback_split(
     f, backend::AutoReverseChainRules, x, ::NoPullbackExtras
 )
@@ -35,13 +28,4 @@ function DI.value_and_pullback_split(
     return y, pullbackfunc
 end
 
-function DI.value_and_pullback!!_split(
-    f, backend::AutoReverseChainRules, x, ::NoPullbackExtras
-)
-    rc = ruleconfig(backend)
-    y, pullback = rrule_via_ad(rc, f, x)
-    pullbackfunc!!(_dx, dy) = last(pullback(dy))
-    return y, pullbackfunc!!
-end
-
 end

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_allocating.jl

@@ -24,8 +24,25 @@ function DI.value_and_pullback(
     return y, new_dx
 end
 
+function DI.value_and_pullback(
+    f, backend::AutoReverseEnzyme, x::AbstractArray, dy, extras::NoPullbackExtras
+)
+    dx = similar(x)
+    return DI.value_and_pullback!!(f, dx, backend, x, dy, extras)
+end
+
+function DI.pullback(f, backend::AutoReverseEnzyme, x, dy, extras::NoPullbackExtras)
+    return DI.value_and_pullback(f, backend, x, dy, extras)[2]
+end
+
 ### In-place
 
+function DI.value_and_pullback!!(
+    f, _dx, backend::AutoReverseEnzyme, x::Number, dy, extras::NoPullbackExtras
+)
+    return DI.value_and_pullback(f, backend, x, dy, extras)
+end
+
 function DI.value_and_pullback!!(
     f, dx, ::AutoReverseEnzyme, x::AbstractArray, dy::Number, ::NoPullbackExtras
 )
@@ -48,11 +65,26 @@ function DI.value_and_pullback!!(
     return y, dx_sametype
 end
 
-function DI.value_and_pullback(
-    f, backend::AutoReverseEnzyme, x::AbstractArray, dy, extras::NoPullbackExtras
+function DI.pullback!!(f, dx, backend::AutoReverseEnzyme, x, dy, extras::NoPullbackExtras)
+    return DI.value_and_pullback!!(f, dx, backend, x, dy, extras)[2]
+end
+
+### Closure
+
+function DI.value_and_pullback_split(
+    f, backend::AutoReverseEnzyme, x, extras::NoPullbackExtras
 )
-    dx = similar(x)
-    return DI.value_and_pullback!!(f, dx, backend, x, dy, extras)
+    y = f(x)
+    pullbackfunc(dy) = DI.pullback(f, backend, x, dy, extras)
+    return y, pullbackfunc
+end
+
+function DI.value_and_pullback!!_split(
+    f, backend::AutoReverseEnzyme, x, extras::NoPullbackExtras
+)
+    y = f(x)
+    pullbackfunc!!(dx, dy) = DI.pullback!!(f, dx, backend, x, dy, extras)
+    return y, pullbackfunc!!
 end
 
 ## Gradient

DifferentiationInterface/ext/DifferentiationInterfaceReverseDiffExt/allocating.jl

@@ -2,79 +2,40 @@
 
 DI.prepare_pullback(f, ::AnyAutoReverseDiff, x) = NoPullbackExtras()
 
-function DI.value_and_pullback!!(
-    f,
-    dx::AbstractArray,
-    backend::AnyAutoReverseDiff,
-    x::AbstractArray,
-    dy,
-    extras::NoPullbackExtras,
-)
-    return f(x), DI.pullback!!(f, dx, backend, x, dy, extras)
-end
-
-function DI.value_and_pullback(
-    f, backend::AnyAutoReverseDiff, x::AbstractArray, dy, extras::NoPullbackExtras
-)
-    return f(x), DI.pullback(f, backend, x, dy, extras)
-end
-
-### Number out
-
-function DI.pullback!!(
-    f,
-    dx::AbstractArray,
-    ::AnyAutoReverseDiff,
-    x::AbstractArray,
-    dy::Number,
-    ::NoPullbackExtras,
+function DI.value_and_pullback_split(
+    f, ::AnyAutoReverseDiff, x::AbstractArray, ::NoPullbackExtras
 )
-    dx = gradient!(dx, f, x)
-    dx .*= dy
-    return dx
-end
-
-function DI.pullback(
-    f, ::AnyAutoReverseDiff, x::AbstractArray, dy::Number, ::NoPullbackExtras
-)
-    dx = gradient(f, x)
-    dx .*= dy
-    return dx
-end
-
-### Array out
-
-function DI.pullback!!(
-    f,
-    dx::AbstractArray,
-    ::AnyAutoReverseDiff,
-    x::AbstractArray,
-    dy::AbstractArray,
-    ::NoPullbackExtras,
-)
-    dotproduct_closure(x) = dot(f(x), dy)
-    dx = gradient!(dx, dotproduct_closure, x)
-    return dx
+    y = f(x)
+    pullbackfunc = if y isa Number
+        dy -> dy .* gradient(f, x)
+    elseif y isa AbstractArray
+        dy -> gradient(z -> dot(f(z), dy), x)
+    end
+    return y, pullbackfunc
 end
 
-function DI.pullback(
-    f, ::AnyAutoReverseDiff, x::AbstractArray, dy::AbstractArray, extras::NoPullbackExtras
+function DI.value_and_pullback!!_split(
+    f, ::AnyAutoReverseDiff, x::AbstractArray, ::NoPullbackExtras
 )
-    dotproduct_closure(x) = dot(f(x), dy)
-    dx = gradient(dotproduct_closure, x)
-    return dx
+    y = f(x)
+    pullbackfunc!! = if y isa Number
+        (dx, dy) -> begin
+            dx = gradient!(dx, f, x)
+            dx .*= dy
+        end
+    elseif y isa AbstractArray
+        (dx, dy) -> gradient!(dx, z -> dot(f(z), dy), x)
+    end
+    return y, pullbackfunc!!
 end
 
-### Number in, not supported
-
-function DI.value_and_pullback(
-    f, backend::AnyAutoReverseDiff, x::Number, dy, ::NoPullbackExtras
+function DI.value_and_pullback_split(
+    f, backend::AnyAutoReverseDiff, x::Number, ::NoPullbackExtras
 )
     x_array = [x]
     f_array = f ∘ only
-    new_extras = DI.prepare_pullback(f_array, backend, x_array)
-    y, dx_array = DI.value_and_pullback(f_array, backend, x_array, dy, new_extras)
-    return y, only(dx_array)
+    y, pullbackfunc = DI.value_and_pullback_split(f_array, backend, x_array)
+    return y, only ∘ pullbackfunc
 end
 
 ## Gradient

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/DifferentiationInterfaceTapirExt.jl

@@ -9,6 +9,7 @@ using Tapir:
     NoTangent,
     build_rrule,
     increment!!,
+    primal,
     set_to_zero!!,
     tangent,
     tangent_type,

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/allocating.jl

@@ -23,3 +23,48 @@ function DI.value_and_pullback!!(
     )
     return new_y, new_dx
 end
+
+function DI.pullback(f, backend::AutoTapir, x, dy, extras::TapirAllocatingPullbackExtras)
+    return DI.value_and_pullback(f, backend, x, dy, extras)[2]
+end
+
+function DI.pullback!!(
+    f, dx, backend::AutoTapir, x, dy, extras::TapirAllocatingPullbackExtras
+)
+    return DI.value_and_pullback!!(f, dx, backend, x, dy, extras)[2]
+end
+
+#=
+# First try
+
+function DI.value_and_pullback_split(f, ::AutoTapir, x, extras::TapirAllocatingPullbackExtras)
+    tf = zero_tangent(f)
+    tx = zero_tangent(x)
+    out, pb!! = extras.rrule(CoDual(f, tf), CoDual(x, tx))
+    y = copy(primal(out))
+    function pullbackfunc(dy)
+        dy_righttype = convert(tangent_type(typeof(y)), copy(dy))
+        ty = increment!!(tangent(out), dy_righttype)
+        res = pb!!(ty, tf, tx)
+        extras.rrule(CoDual(f, tf), CoDual(x, tx))
+        return last(res)
+    end
+    return y, pullbackfunc
+end
+=#
+
+function DI.value_and_pullback_split(
+    f, backend::AutoTapir, x, extras::TapirAllocatingPullbackExtras
+)
+    y = f(x)
+    pullbackfunc(dy) = DI.pullback(f, backend, x, dy, extras)
+    return y, pullbackfunc
+end
+
+function DI.value_and_pullback!!_split(
+    f, backend::AutoTapir, x, extras::TapirAllocatingPullbackExtras
+)
+    y = f(x)
+    pullbackfunc!!(dx, dy) = DI.pullback!!(f, dx, backend, x, dy, extras)
+    return y, pullbackfunc!!
+end

DifferentiationInterface/ext/DifferentiationInterfaceTrackerExt/DifferentiationInterfaceTrackerExt.jl

@@ -11,23 +11,12 @@ DI.supports_mutation(::AutoTracker) = DI.MutationNotSupported()
 
 DI.prepare_pullback(f, ::AutoTracker, x) = NoPullbackExtras()
 
-function DI.value_and_pullback(f, ::AutoTracker, x, dy, ::NoPullbackExtras)
-    y, back = forward(f, x)
-    return y, data(only(back(dy)))
-end
-
 function DI.value_and_pullback_split(f, ::AutoTracker, x, ::NoPullbackExtras)
     y, back = forward(f, x)
     pullbackfunc(dy) = data(only(back(dy)))
     return y, pullbackfunc
 end
 
-function DI.value_and_pullback!!_split(f, ::AutoTracker, x, ::NoPullbackExtras)
-    y, back = forward(f, x)
-    pullbackfunc!!(_dx, dy) = data(only(back(dy)))
-    return y, pullbackfunc!!
-end
-
 ## Gradient
 
 DI.prepare_gradient(f, ::AutoTracker, x) = NoGradientExtras()

DifferentiationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl

@@ -16,24 +16,12 @@ DI.supports_mutation(::AnyAutoZygote) = DI.MutationNotSupported()
 
 DI.prepare_pullback(f, ::AnyAutoZygote, x) = NoPullbackExtras()
 
-function DI.value_and_pullback(f, ::AnyAutoZygote, x, dy, ::NoPullbackExtras)
-    y, back = pullback(f, x)
-    dx = only(back(dy))
-    return y, dx
-end
-
 function DI.value_and_pullback_split(f, ::AnyAutoZygote, x, ::NoPullbackExtras)
     y, back = pullback(f, x)
     pullbackfunc(dy) = only(back(dy))
     return y, pullbackfunc
 end
 
-function DI.value_and_pullback!!_split(f, ::AnyAutoZygote, x, ::NoPullbackExtras)
-    y, back = pullback(f, x)
-    pullbackfunc!!(_dx, dy) = only(back(dy))
-    return y, pullbackfunc!!
-end
-
 ## Gradient
 
 DI.prepare_gradient(f, ::AnyAutoZygote, x) = NoGradientExtras()

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -130,20 +130,4 @@ export prepare_second_derivative, prepare_hvp, prepare_hessian
 
 export check_available, check_mutation, check_hessian
 
-function __init__()
-    Base.Experimental.register_error_hint(StackOverflowError) do io, exc
-        print(
-            io,
-            """\n
-            HINT: One of DifferentiationInterface's functions might be missing a method, which would trigger an endless loop of `pullback` calling `pushforward` and vice-versa.
-            Some possible fixes:
-            - switch to another backend
-            - if you don't want to switch, load the package extension corresponding to your backend
-            - if your backend is already loaded, define the primitive operator for the right combination of argument types
-            """,
-        )
-        return nothing
-    end
-end
-
 end # module