Make value_and_pullback the primitive (#157)

Author: gdalle

DifferentiationInterface/docs/src/backends.md

@@ -18,7 +18,7 @@ function all_backends()
         AutoFiniteDiff(),
         AutoFiniteDifferences(FiniteDifferences.central_fdm(3, 1)),
         AutoForwardDiff(),
-        AutoPolyesterForwardDiff(; chunksize=2),
+        AutoPolyesterForwardDiff(; chunksize=1),
         AutoReverseDiff(),
         AutoTapir(),
         AutoTracker(),

DifferentiationInterface/docs/src/design.md

@@ -17,12 +17,5 @@ Advanced users are welcome to code more backends and submit pull requests!
 
 ## Fallback call structure
 
-For simplicity, we remove `value_` in the operator names below.
-
-!!! note "Edge labels"
-
-    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.
-
 !!! warning
-    This is still in flux, come back later!
+    This is still in flux, come back later!

DifferentiationInterface/docs/src/overview.md

@@ -7,11 +7,11 @@ Most backends have custom implementations, which we reuse if possible.
 
 We choose the following terminology for the high-level operators we provide:
 
-| operator             | input  `x`      | output   `y`    | result type      | result shape             |
-| -------------------- | --------------- | --------------- | ---------------- | ------------------------ |
-| [`derivative`](@ref) | `Number`        | `Any`           | same as `y`      | `size(y)`                |
-| [`gradient`](@ref)   | `Any`           | `Number`        | same as `x`      | `size(x)`                |
-| [`jacobian`](@ref)   | `AbstractArray` | `AbstractArray` | `AbstractMatrix` | `(length(y), length(x))` |
+| operator             | input  `x`      | output   `y`                | result type      | result shape             |
+| -------------------- | --------------- | --------------------------- | ---------------- | ------------------------ |
+| [`derivative`](@ref) | `Number`        | `Number` or `AbstractArray` | same as `y`      | `size(y)`                |
+| [`gradient`](@ref)   | `AbstractArray` | `Number`                    | same as `x`      | `size(x)`                |
+| [`jacobian`](@ref)   | `AbstractArray` | `AbstractArray`             | `AbstractMatrix` | `(length(y), length(x))` |
 
 They are all based on the following low-level operators:
 
@@ -58,11 +58,11 @@ You can either pick a single backend to do all the work, or combine an "outer" b
 
 The available operators are similar to first-order ones:
 
-| operator                    | input  `x`      | output   `y` | result type      | result shape             |
-| --------------------------- | --------------- | ------------ | ---------------- | ------------------------ |
-| [`second_derivative`](@ref) | `Number`        | `Any`        | same as `y`      | `size(y)`                |
-| [`hvp`](@ref)               | `Any`           | `Number`     | same as `x`      | `size(x)`                |
-| [`hessian`](@ref)           | `AbstractArray` | `Number`     | `AbstractMatrix` | `(length(x), length(x))` |
+| operator                    | input  `x`      | output   `y`                | result type      | result shape             |
+| --------------------------- | --------------- | --------------------------- | ---------------- | ------------------------ |
+| [`second_derivative`](@ref) | `Number`        | `Number` or `AbstractArray` | same as `y`      | `size(y)`                |
+| [`hvp`](@ref)               | `AbstractArray` | `Number`                    | same as `x`      | `size(x)`                |
+| [`hessian`](@ref)           | `AbstractArray` | `Number`                    | `AbstractMatrix` | `(length(x), length(x))` |
 
 We only define two variants for now:
 
@@ -94,12 +94,9 @@ This is a backend-specific procedure, but we expose a common syntax to achieve i
 If you run `prepare_operator(backend, f, x)`, it will create an object called `extras` containing the necessary information to speed up `operator` and its variants.
 This information is specific to `backend` and `f`, as well as the _type and size_ of the input `x`, but it should work with different _values_ of `x`.
 
-You can then call `operator(backend, f, similar_x, extras)`, which should be faster than `operator(backend, f, similar_x)`.
+You can then call `operator(backend, f, x2, extras)`, which should be faster than `operator(f, backend, x2)`.
 This is especially worth it if you plan to call `operator` several times in similar settings: you can think of it as a warm up.
 
-By default, all the preparation functions return `nothing`.
-We do not make any guarantees on their implementation for each backend, or on the performance gains that can be expected.
-
 !!! warning
     For `SecondOrder` backends, the inner differentiation cannot be prepared at the moment, only the outer one is.
 
@@ -123,12 +120,13 @@ This means the Hessian is obtained as the sparse Jacobian of the gradient.
 
 ### Split reverse mode
 
-Many reverse mode AD backends expose a "split" option, which runs only the forward sweep, and encapsulates the reverse sweep in a closure.
-We make this available for allocating functions only, with the following operators:
+Some reverse mode AD backends expose a "split" option, which runs only the forward sweep, and encapsulates the reverse sweep in a closure.
+We make this available for all backends with the following operators:
 
-| out-of-place                       | in-place (or not)                    |
-| ---------------------------------- | ------------------------------------ |
-| [`value_and_pullback_split`](@ref) | [`value_and_pullback!!_split`](@ref) |
+|                      | out-of-place                       | in-place (or not)                      |
+| -------------------- | ---------------------------------- | -------------------------------------- |
+| allocating functions | [`value_and_pullback_split`](@ref) | [`value_and_pullback!!_split`](@ref)   |
+| mutating functions   | -                                  | [`value_and_pullback!!_split!!`](@ref) |
 
 !!! danger
     Split reverse mode is still experimental, use at your own risk.

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -17,7 +17,7 @@ DI.mode(::AutoReverseChainRules) = ADTypes.AbstractReverseMode
 
 ## Pullback
 
-DI.prepare_pullback(f, ::AutoForwardChainRules, x) = NoPullbackExtras()
+DI.prepare_pullback(f, ::AutoReverseChainRules, x) = NoPullbackExtras()
 
 function DI.value_and_pullback_split(
     f, backend::AutoReverseChainRules, x, ::NoPullbackExtras
@@ -28,4 +28,11 @@ function DI.value_and_pullback_split(
     return y, pullbackfunc
 end
 
+function DI.value_and_pullback(
+    f, backend::AutoReverseChainRules, x, dy, extras::NoPullbackExtras
+)
+    y, pullbackfunc = DI.value_and_pullback_split(f, backend, x, extras)
+    return y, pullbackfunc(dy)
+end
+
 end

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/reverse_allocating.jl

@@ -69,24 +69,6 @@ function DI.pullback!!(f, dx, backend::AutoReverseEnzyme, x, dy, extras::NoPullb
     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
-)
-    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
 
 DI.prepare_gradient(f, ::AutoReverseEnzyme, x) = NoGradientExtras()

DifferentiationInterface/ext/DifferentiationInterfaceReverseDiffExt/allocating.jl

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

DifferentiationInterface/ext/DifferentiationInterfaceTapirExt/allocating.jl

@@ -38,19 +38,3 @@ function DI.pullback!!(
 )
     return DI.value_and_pullback!!(f, dx, backend, x, dy, extras)[2]
 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

@@ -17,6 +17,11 @@ function DI.value_and_pullback_split(f, ::AutoTracker, x, ::NoPullbackExtras)
     return y, pullbackfunc
 end
 
+function DI.value_and_pullback(f, backend::AutoTracker, x, dy, extras::NoPullbackExtras)
+    y, pullbackfunc = DI.value_and_pullback_split(f, backend, x, extras)
+    return y, pullbackfunc(dy)
+end
+
 ## Gradient
 
 DI.prepare_gradient(f, ::AutoTracker, x) = NoGradientExtras()

DifferentiationInterface/ext/DifferentiationInterfaceZygoteExt/DifferentiationInterfaceZygoteExt.jl

@@ -22,6 +22,11 @@ function DI.value_and_pullback_split(f, ::AnyAutoZygote, x, ::NoPullbackExtras)
     return y, pullbackfunc
 end
 
+function DI.value_and_pullback(f, backend::AnyAutoZygote, x, dy, extras::NoPullbackExtras)
+    y, pullbackfunc = DI.value_and_pullback_split(f, backend, x, extras)
+    return y, pullbackfunc(dy)
+end
+
 ## Gradient
 
 DI.prepare_gradient(f, ::AnyAutoZygote, x) = NoGradientExtras()

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -107,7 +107,7 @@ export SecondOrder
 
 export value_and_pushforward!!, value_and_pushforward
 export value_and_pullback!!, value_and_pullback
-export value_and_pullback!!_split, value_and_pullback_split
+export value_and_pullback!!_split, value_and_pullback_split, value_and_pullback!!_split!!
 
 export value_and_derivative!!, value_and_derivative
 export value_and_gradient!!, value_and_gradient