Add split reverse mode (#137)

Author: gdalle

DifferentiationInterface/docs/src/overview.md

@@ -120,3 +120,15 @@ This means the Hessian is obtained as the sparse Jacobian of the gradient.
 
 !!! danger
     Sparsity support is still experimental, use at your own risk.
+
+### 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 everyone with the following operators:
+
+| out-of-place                       | in-place (or not)                    |
+| ---------------------------------- | ------------------------------------ |
+| [`value_and_pullback_split`](@ref) | [`value_and_pullback!!_split`](@ref) |
+
+!!! danger
+    Split reverse mode is still experimental, use at your own risk.

DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl

@@ -15,20 +15,6 @@ DI.supports_mutation(::AutoChainRules) = DI.MutationNotSupported()
 DI.mode(::AutoForwardChainRules) = ADTypes.AbstractForwardMode
 DI.mode(::AutoReverseChainRules) = ADTypes.AbstractReverseMode
 
-## Pushforward (unused)
-
-#=
-DI.prepare_pushforward(f, ::AutoForwardChainRules, x) = NoPushforwardExtras()
-
-function DI.value_and_pushforward(
-    f, backend::AutoForwardChainRules, x, dx, ::NoPushforwardExtras
-)
-    rc = ruleconfig(backend)
-    y, new_dy = frule_via_ad(rc, (NoTangent(), dx), f, x)
-    return y, new_dy
-end
-=#
-
 ## Pullback
 
 DI.prepare_pullback(f, ::AutoForwardChainRules, x) = NoPullbackExtras()
@@ -40,4 +26,22 @@ function DI.value_and_pullback(f, backend::AutoReverseChainRules, x, dy, ::NoPul
     return y, new_dx
 end
 
+function DI.value_and_pullback_split(
+    f, backend::AutoReverseChainRules, x, ::NoPullbackExtras
+)
+    rc = ruleconfig(backend)
+    y, pullback = rrule_via_ad(rc, f, x)
+    pullbackfunc(dy) = last(pullback(dy))
+    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/DifferentiationInterfaceTrackerExt/DifferentiationInterfaceTrackerExt.jl

@@ -16,6 +16,18 @@ function DI.value_and_pullback(f, ::AutoTracker, x, dy, ::NoPullbackExtras)
     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

@@ -22,6 +22,18 @@ function DI.value_and_pullback(f, ::AnyAutoZygote, x, dy, ::NoPullbackExtras)
     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

@@ -107,6 +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_derivative!!, value_and_derivative
 export value_and_gradient!!, value_and_gradient
@@ -130,34 +131,18 @@ 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, argtypes, kwargs
-        f_name = string(exc.f)
-        if (
-            f_name == "mode" ||
-            contains(f_name, "pushforward") ||
-            contains(f_name, "pullback") ||
-            contains(f_name, "derivative") ||
-            contains(f_name, "gradient") ||
-            contains(f_name, "jacobian") ||
-            contains(f_name, "hvp") ||
-            contains(f_name, "hessian")
+    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
+            """,
         )
-            for T in argtypes
-                if T <: AbstractADType
-                    print(
-                        io,
-                        """\n
-                        HINT: One of DifferentiationInterface's functions is missing a method, which causes 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 backend `$T`
-                        - if the package is already loaded, define the method `$f_name` for the right combination of argument types
-                        """,
-                    )
-                    return nothing
-                end
-            end
-        end
+        return nothing
     end
 end
 

DifferentiationInterface/src/jacobian.jl

@@ -64,9 +64,7 @@ function value_and_jacobian_aux(
     y = f(x)
     jac = stack(CartesianIndices(x); dims=2) do j
         dx_j = basis(backend, x, j)
-        jac_col_j = last(
-            value_and_pushforward(f, backend, x, dx_j, extras.pushforward_extras)
-        )
+        jac_col_j = pushforward(f, backend, x, dx_j, extras.pushforward_extras)
         vec(jac_col_j)
     end
     return y, jac
@@ -75,10 +73,10 @@ end
 function value_and_jacobian_aux(
     f, backend, x::AbstractArray, extras::PullbackJacobianExtras
 )
-    y = f(x)
+    y, pullbackfunc = value_and_pullback_split(f, backend, x, extras.pullback_extras)
     jac = stack(CartesianIndices(y); dims=1) do i
         dy_i = basis(backend, y, i)
-        jac_row_i = last(value_and_pullback(f, backend, x, dy_i, extras.pullback_extras))
+        jac_row_i = pullbackfunc(dy_i)
         vec(jac_row_i)
     end
     return y, jac
@@ -116,13 +114,11 @@ end
 function value_and_jacobian_aux!!(
     f, jac::AbstractMatrix, backend, x::AbstractArray, extras::PullbackJacobianExtras
 )
-    y = f(x)
+    y, pullbackfunc!! = value_and_pullback!!_split(f, backend, x, extras.pullback_extras)
     for (k, i) in enumerate(CartesianIndices(y))
         dy_i = basis(backend, y, i)
         jac_row_i_old = reshape(view(jac, k, :), size(x))
-        jac_row_i_new = pullback!!(
-            f, jac_row_i_old, backend, x, dy_i, extras.pullback_extras
-        )
+        jac_row_i_new = pullbackfunc!!(jac_row_i_old, dy_i)
         # this allocates
         copyto!(jac_row_i_old, jac_row_i_new)
     end
@@ -175,7 +171,6 @@ function value_and_jacobian_aux!!(
     x::AbstractArray,
     extras::PushforwardJacobianExtras,
 )
-    f!(y, x)
     for (k, j) in enumerate(CartesianIndices(x))
         dx_j = basis(backend, x, j)
         jac_col_j_old = reshape(view(jac, :, k), size(y))
@@ -198,15 +193,13 @@ function value_and_jacobian_aux!!(
     x::AbstractArray,
     extras::PullbackJacobianExtras,
 )
-    f!(y, x)
+    y, pullbackfunc!! = value_and_pullback!!_split(
+        f!, y, backend, x, extras.pullback_extras
+    )
     for (k, i) in enumerate(CartesianIndices(y))
         dy_i = basis(backend, y, i)
         jac_row_i_old = reshape(view(jac, k, :), size(x))
-        jac_row_i_new = last(
-            value_and_pullback!!(
-                f!, y, jac_row_i_old, backend, x, dy_i, extras.pullback_extras
-            ),
-        )
+        jac_row_i_new = pullbackfunc!!(y, jac_row_i_old, dy_i)
         # this allocates
         copyto!(jac_row_i_old, jac_row_i_new)
     end

DifferentiationInterface/src/pullback.jl

@@ -124,3 +124,60 @@ function value_and_pullback!!(
     end
     return y, dx
 end
+
+## Closure
+
+"""
+    value_and_pullback_split(f, backend, x, [extras])
+
+Apply split reverse mode autodiff.
+
+Returns a tuple `(y, pullbackfunc)` where the second element is a function (closure) with the following signature:
+
+    pullbackfunc(dy) -> dx
+"""
+function value_and_pullback_split(
+    f, backend::AbstractADType, x, extras::PullbackExtras=prepare_pullback(f, backend, x)
+)
+    pullbackfunc(dy) = pullback(f, backend, x, dy, extras)
+    return f(x), pullbackfunc
+end
+
+"""
+    value_and_pullback!!_split(f, backend, x, [extras])
+
+Apply split reverse mode autodiff.
+
+Returns a tuple `(y, pullbackfunc!!)` where the second element is a function (closure) with the following signature:
+
+    pullbackfunc!!(dx, dy) -> dx
+"""
+function value_and_pullback!!_split(
+    f, backend::AbstractADType, x, extras::PullbackExtras=prepare_pullback(f, backend, x)
+)
+    pullbackfunc!!(dx, dy) = pullback!!(f, dx, backend, x, dy, extras)
+    return f(x), pullbackfunc!!
+end
+
+"""
+    value_and_pullback!!_split(f!, y, backend, x, [extras])
+
+Apply split reverse mode autodiff.
+
+Returns a tuple `(y, pullbackfunc!!)` where the second element is a function (closure) with the following signature:
+
+    pullbackfunc!!(y, dx, dy) -> dx
+"""
+function value_and_pullback!!_split(
+    f!,
+    y,
+    backend::AbstractADType,
+    x,
+    extras::PullbackExtras=prepare_pullback(f!, backend, y, x),
+)
+    function pullbackfunc!!(y, dx, dy)
+        return value_and_pullback!!(f!, y, dx, backend, x, dy, extras)[2]
+    end
+    f!(y, x)
+    return y, pullbackfunc!!
+end

DifferentiationInterfaceTest/src/tests/benchmark.jl

@@ -37,7 +37,7 @@ end
 function record!(
     data::Vector{BenchmarkDataRow},
     backend::AbstractADType,
-    operator::Function,
+    operator,
     scenario::AbstractScenario,
     bench,
 )
@@ -98,10 +98,13 @@ function run_benchmark!(
 )
     (; f, x, y, dy) = deepcopy(scen)
     extras = prepare_pullback(f, ba, x)
+    _, pullbackfunc!! = value_and_pullback!!_split(f, ba, x, extras)
     bench1 = @be mysimilar(x) value_and_pullback!!(f, _, ba, x, dy, extras)
     bench2 = @be mysimilar(x) pullback!!(f, _, ba, x, dy, extras)
+    bench3 = @be mysimilar(x) pullbackfunc!!(_, dy)
     record!(data, ba, value_and_pullback!!, scen, bench1)
     record!(data, ba, pullback!!, scen, bench2)
+    record!(data, ba, pullbackfunc!!, scen, bench3)
     return nothing
 end
 
@@ -111,10 +114,13 @@ function run_benchmark!(
     (; f, x, y, dy) = deepcopy(scen)
     f! = f
     extras = prepare_pullback(f!, ba, y, x)
+    _, pullbackfunc!! = value_and_pullback!!_split(f!, y, ba, x, extras)
     bench1 = @be (mysimilar(y), mysimilar(x)) value_and_pullback!!(
         f!, _[1], _[2], ba, x, dy, extras
     )
+    bench2 = @be (mysimilar(y), mysimilar(x)) pullbackfunc!!(_[1], _[2], dy)
     record!(data, ba, value_and_pullback!!, scen, bench1)
+    record!(data, ba, pullbackfunc!!, scen, bench2)
     return nothing
 end
 

DifferentiationInterfaceTest/src/tests/correctness.jl

@@ -109,16 +109,25 @@ function test_correctness(
     dx3 = pullback(f, ba, x, dy, extras)
     dx4 = pullback!!(f, mysimilar(x), ba, x, dy, extras)
 
+    y5, pullbackfunc = value_and_pullback_split(f, ba, x, extras)
+    dx5 = pullbackfunc(dy)
+    y6, pullbackfunc!! = value_and_pullback!!_split(f, ba, x, extras)
+    dx6 = pullbackfunc!!(mysimilar(x), dy)
+
     let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y1 ≈ y
             @test y2 ≈ y
+            @test y5 ≈ y
+            @test y6 ≈ y
         end
         @testset "Cotangent value" begin
             @test dx1 ≈ dx_true
             @test dx2 ≈ dx_true
             @test dx3 ≈ dx_true
             @test dx4 ≈ dx_true
+            @test dx5 ≈ dx_true
+            @test dx6 ≈ dx_true
         end
     end
     test_scen_intact(new_scen, scen)
@@ -145,13 +154,20 @@ function test_correctness(
     y10 = mysimilar(y)
     y1, dx1 = value_and_pullback!!(f!, y10, mysimilar(x), ba, x, dy, extras)
 
+    y20 = mysimilar(y)
+    y2, pullbackfunc!! = value_and_pullback!!_split(f!, y20, ba, x, extras)
+    dx2 = pullbackfunc!!(y20, mysimilar(x), dy)
+
     let (≈)(x, y) = isapprox(x, y; atol, rtol)
         @testset "Primal value" begin
             @test y10 ≈ y
+            @test y20 ≈ y
             @test y1 ≈ y
+            @test y2 ≈ y
         end
         @testset "Cotangent value" begin
             @test dx1 ≈ dx_true
+            @test dx2 ≈ dx_true
         end
     end
     test_scen_intact(new_scen, scen)

DifferentiationInterfaceTest/src/tests/type_stability.jl

@@ -32,9 +32,14 @@ function test_jet(ba::AbstractADType, scen::PullbackScenario{false};)
     extras = prepare_pullback(f, ba, x)
     dx_in = mysimilar(x)
 
+    _, pullbackfunc!! = value_and_pullback!!_split(f, ba, x, extras)
+    _, pullbackfunc = value_and_pullback_split(f, ba, x, extras)
+
     if Bool(pullback_performance(ba))
         @test_opt value_and_pullback!!(f, dx_in, ba, x, dy, extras)
         @test_opt value_and_pullback(f, ba, x, dy, extras)
+        @test_opt pullbackfunc!!(dx_in, dy)
+        @test_opt pullbackfunc(dy)
     end
     return nothing
 end
@@ -46,8 +51,11 @@ function test_jet(ba::AbstractADType, scen::PullbackScenario{true};)
     y_in = mysimilar(y)
     dx_in = mysimilar(x)
 
+    _, pullbackfunc!! = value_and_pullback!!_split(f!, y, ba, x, extras)
+
     if Bool(pullback_performance(ba))
         @test_opt value_and_pullback!!(f!, y_in, dx_in, ba, x, dy, extras)
+        @test_opt pullbackfunc!!(y_in, dx_in, dy)
     end
     return nothing
 end