feat: Add reverse/forward counterparts

DifferentiationInterface/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.18...main)
 
+### Added
+
+- Add `forward_counterpart` and `reverse_counterpart` to retrieve the forward- or reverse-mode counterpart of a backend ([#1025](https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/1025))
+
 ## [0.7.18](https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.7.17...DifferentiationInterface-v0.7.18)
 
 ### Added

DifferentiationInterface/docs/src/api.md

@@ -123,6 +123,8 @@ DifferentiationInterface.inner
 
 ```@docs
 DifferentiateWith
+DifferentiationInterface.forward_counterpart
+DifferentiationInterface.reverse_counterpart
 ```
 
 ### Sparsity tools

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl

@@ -47,6 +47,7 @@ using Enzyme:
 DI.check_available(::AutoEnzyme) = true
 
 include("utils.jl")
+include("counterparts.jl")
 
 include("forward_onearg.jl")
 include("forward_twoarg.jl")

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/counterparts.jl

@@ -0,0 +1,42 @@
+## Backend counterparts
+
+# Flip the Enzyme mode while preserving its other type parameters (and the function
+# annotation `A`).
+
+function DI.forward_counterpart(
+        ::AutoEnzyme{
+            <:ReverseMode{
+                ReturnPrimal, RuntimeActivity, StrongZero, ABI, Holomorphic, ErrIfFuncWritten,
+            },
+            A,
+        },
+    ) where {ReturnPrimal, RuntimeActivity, StrongZero, ABI, Holomorphic, ErrIfFuncWritten, A}
+    mode = ForwardMode{ReturnPrimal, ABI, ErrIfFuncWritten, RuntimeActivity, StrongZero}()
+    return AutoEnzyme(; mode, function_annotation = A)
+end
+
+function DI.forward_counterpart(
+        ::AutoEnzyme{
+            <:ReverseModeSplit{
+                ReturnPrimal, ReturnShadow, RuntimeActivity, StrongZero, Width,
+                ModifiedBetween, ABI, Holomorphic, ErrIfFuncWritten, ShadowInit,
+            },
+            A,
+        },
+    ) where {
+        ReturnPrimal, ReturnShadow, RuntimeActivity, StrongZero, Width,
+        ModifiedBetween, ABI, Holomorphic, ErrIfFuncWritten, ShadowInit, A,
+    }
+    mode = ForwardMode{ReturnPrimal, ABI, ErrIfFuncWritten, RuntimeActivity, StrongZero}()
+    return AutoEnzyme(; mode, function_annotation = A)
+end
+
+function DI.reverse_counterpart(
+        ::AutoEnzyme{
+            <:ForwardMode{ReturnPrimal, ABI, ErrIfFuncWritten, RuntimeActivity, StrongZero},
+            A,
+        },
+    ) where {ReturnPrimal, ABI, ErrIfFuncWritten, RuntimeActivity, StrongZero, A}
+    mode = ReverseMode{ReturnPrimal, RuntimeActivity, StrongZero, ABI, false, ErrIfFuncWritten}()
+    return AutoEnzyme(; mode, function_annotation = A)
+end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -39,6 +39,7 @@ DI.inner_preparation_behavior(::AutoMooncakeForward) = DI.PrepareInnerSimple()
 @inline new_friendly_tangents() = isdefined(Mooncake, :FriendlyTangentCache)
 
 include("utils.jl")
+include("counterparts.jl")
 include("onearg.jl")
 include("twoarg.jl")
 include("forward_onearg.jl")

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/counterparts.jl

@@ -0,0 +1,4 @@
+## Backend counterparts
+
+DI.forward_counterpart(backend::AutoMooncake) = AutoMooncakeForward(; config = backend.config)
+DI.reverse_counterpart(backend::AutoMooncakeForward) = AutoMooncake(; config = backend.config)

DifferentiationInterface/src/DifferentiationInterface.jl

@@ -49,6 +49,7 @@ include("utils/traits.jl")
 include("utils/basis.jl")
 include("utils/batchsize.jl")
 include("utils/check.jl")
+include("utils/counterparts.jl")
 include("utils/errors.jl")
 include("utils/linalg.jl")
 include("utils/sparse.jl")
@@ -132,6 +133,7 @@ export AutoSparse
 @public inner, outer
 @public AutoForwardFromPrimitive, AutoReverseFromPrimitive
 @public Prep
+@public forward_counterpart, reverse_counterpart
 
 include("init.jl")
 

DifferentiationInterface/src/misc/from_primitive.jl

@@ -270,3 +270,16 @@ function value_and_pullback!(
         f!, y, tx, prep.pullback_prep, backend.backend, x, ty, contexts...
     )
 end
+
+## Counterparts
+
+# The counterpart of a `FromPrimitive` wrapper swaps the primitive (pushforward or
+# pullback), applying the counterpart to the wrapped backend as well.
+
+function forward_counterpart(backend::AutoReverseFromPrimitive)
+    return AutoForwardFromPrimitive(forward_counterpart(backend.backend))
+end
+
+function reverse_counterpart(backend::AutoForwardFromPrimitive)
+    return AutoReverseFromPrimitive(reverse_counterpart(backend.backend))
+end

DifferentiationInterface/src/misc/zero_backends.jl

@@ -201,3 +201,10 @@ function value_and_pullback!(
     end
     return y, tx
 end
+
+## Counterparts
+
+# The zero backends are each other's counterpart, since they both return zero derivatives.
+
+forward_counterpart(::AutoZeroReverse) = AutoZeroForward()
+reverse_counterpart(::AutoZeroForward) = AutoZeroReverse()

DifferentiationInterface/src/utils/counterparts.jl

@@ -0,0 +1,85 @@
+"""
+    forward_counterpart(backend)
+
+Return a forward-mode counterpart of `backend`.
+
+If `backend` has a dedicated forward-mode counterpart (e.g. `AutoMooncake` has
+`AutoMooncakeForward`), it is returned. Else, `backend` itself is returned (DI allows
+reverse-mode backends to execute pushforwards), with a warning in the case it is
+reverse-mode only.
+"""
+function forward_counterpart(backend::AbstractADType)
+    if !(mode(backend) isa Union{ForwardMode, ForwardOrReverseMode, SymbolicMode})
+        @warn "The forward-mode counterpart of `$backend` is itself, returning it unchanged." maxlog =
+            1
+    end
+    return backend
+end
+
+"""
+    reverse_counterpart(backend)
+
+Return a reverse-mode counterpart of `backend`.
+
+If `backend` has a dedicated reverse-mode counterpart (e.g. `AutoMooncakeForward` has
+`AutoMooncake`), it is returned. Else, `backend` itself is returned (DI allows
+forward-mode backends to execute pullbacks), with a warning in the case it is
+forward-mode only.
+"""
+function reverse_counterpart(backend::AbstractADType)
+    if !(mode(backend) isa Union{ReverseMode, ForwardOrReverseMode, SymbolicMode})
+        @warn "The reverse-mode counterpart of `$backend` is itself, returning it unchanged." maxlog =
+            1
+    end
+    return backend
+end
+
+## Wrapper backends
+
+# The counterpart of `AutoSparse` acts on the dense backend and preserves the sparsity
+# detection and coloring machinery.
+
+function forward_counterpart(backend::AutoSparse)
+    return AutoSparse(
+        forward_counterpart(dense_ad(backend));
+        sparsity_detector = backend.sparsity_detector,
+        coloring_algorithm = backend.coloring_algorithm,
+    )
+end
+
+function reverse_counterpart(backend::AutoSparse)
+    return AutoSparse(
+        reverse_counterpart(dense_ad(backend));
+        sparsity_detector = backend.sparsity_detector,
+        coloring_algorithm = backend.coloring_algorithm,
+    )
+end
+
+# A `MixedMode` backend already combines a forward- and a reverse-mode backend, each used
+# where it works best, so it is its own counterpart in both directions.
+
+forward_counterpart(backend::MixedMode) = backend
+reverse_counterpart(backend::MixedMode) = backend
+
+# For `SecondOrder` there is no meaningful counterpart: flipping the modes of the inner
+# and outer backends changes the nature of the second-order combination.
+
+function forward_counterpart(::SecondOrder)
+    throw(
+        ArgumentError(
+            "`forward_counterpart` is ambiguous for `SecondOrder` backends, apply it to `inner(backend)` and `outer(backend)` separately.",
+        )
+    )
+end
+
+function reverse_counterpart(::SecondOrder)
+    throw(
+        ArgumentError(
+            "`reverse_counterpart` is ambiguous for `SecondOrder` backends, apply it to `inner(backend)` and `outer(backend)` separately.",
+        )
+    )
+end
+
+# Counterparts for `FromPrimitive` wrappers are defined in `misc/from_primitive.jl`, and
+# backend-specific counterparts (e.g., for AutoMooncake and AutoEnzyme) are defined in the
+# corresponding package extensions.

DifferentiationInterface/test/Back/ChainRules/test.jl

@@ -10,6 +10,7 @@ check_no_implicit_imports(DifferentiationInterface)
 for backend in [AutoChainRules(ZygoteRuleConfig())]
     @test check_available(backend)
     @test !check_inplace(backend)
+    test_counterparts(backend)
 end
 
 test_differentiation(

DifferentiationInterface/test/Back/Enzyme/test.jl

@@ -27,9 +27,30 @@ duplicated_backends = [
     @testset "Check $(typeof(backend))" for backend in backends
         @test check_available(backend)
         @test check_inplace(backend)
+        test_counterparts(backend)
     end
 end;
 
+@testset "Counterpart mode attributes" begin
+    # a reverse mode with non-default attributes, to check they survive the flip
+    rev_mode = Enzyme.set_runtime_activity(Enzyme.ReverseWithPrimal)
+    rev = AutoEnzyme(; mode = rev_mode, function_annotation = Enzyme.Const)
+    fwd = DifferentiationInterface.forward_counterpart(rev)
+    @test ADTypes.mode(fwd) isa ADTypes.ForwardMode
+    @test fwd isa AutoEnzyme{<:Any, Enzyme.Const}  # function annotation preserved
+    # ReturnPrimal and RuntimeActivity must carry over (ForwardMode{ReturnPrimal,...,RuntimeActivity,...})
+    @test fwd.mode isa Enzyme.EnzymeCore.ForwardMode{true, <:Any, <:Any, true}
+    # round-trip back to reverse recovers the original mode
+    @test DifferentiationInterface.reverse_counterpart(fwd) === rev
+    # split reverse mode also maps to forward mode, carrying over its attributes
+    split_mode = Enzyme.set_runtime_activity(Enzyme.ReverseSplitWithPrimal)
+    split = AutoEnzyme(; mode = split_mode, function_annotation = Enzyme.Const)
+    fwd_split = DifferentiationInterface.forward_counterpart(split)
+    @test ADTypes.mode(fwd_split) isa ADTypes.ForwardMode
+    @test fwd_split isa AutoEnzyme{<:Any, Enzyme.Const}
+    @test fwd_split.mode isa Enzyme.EnzymeCore.ForwardMode{true, <:Any, <:Any, true}
+end;
+
 @testset "First order" begin
     test_differentiation(
         backends, default_scenarios(); excluded = SECOND_ORDER, logging = LOGGING

DifferentiationInterface/test/Back/FastDifferentiation/test.jl

@@ -12,6 +12,7 @@ check_no_implicit_imports(DifferentiationInterface)
 for backend in [AutoFastDifferentiation(), AutoSparse(AutoFastDifferentiation())]
     @test check_available(backend)
     @test check_inplace(backend)
+    test_counterparts(backend)
 end
 
 test_differentiation(

DifferentiationInterface/test/Back/FiniteDiff/test.jl

@@ -12,6 +12,7 @@ check_no_implicit_imports(DifferentiationInterface)
 for backend in [AutoFiniteDiff()]
     @test check_available(backend)
     @test check_inplace(backend)
+    test_counterparts(backend)
     @test DifferentiationInterface.inner_preparation_behavior(backend) isa
         DifferentiationInterface.PrepareInnerSimple
 end

DifferentiationInterface/test/Back/FiniteDifferences/test.jl

@@ -10,6 +10,7 @@ check_no_implicit_imports(DifferentiationInterface)
 for backend in [AutoFiniteDifferences(; fdm = FiniteDifferences.central_fdm(3, 1))]
     @test check_available(backend)
     @test !check_inplace(backend)
+    test_counterparts(backend)
     @test DifferentiationInterface.inner_preparation_behavior(backend) isa
         DifferentiationInterface.PrepareInnerSimple
 end

DifferentiationInterface/test/Back/ForwardDiff/test.jl

@@ -25,6 +25,7 @@ backends = [
 for backend in backends
     @test check_available(backend)
     @test check_inplace(backend)
+    test_counterparts(backend)
 end
 
 @testset "Dense" begin

DifferentiationInterface/test/Back/GTPSA/test.jl

@@ -11,6 +11,7 @@ check_no_implicit_imports(DifferentiationInterface)
 for backend in [AutoGTPSA()]
     @test check_available(backend)
     @test check_inplace(backend)
+    test_counterparts(backend)
 end
 
 # Test no Descriptor (use context)

DifferentiationInterface/test/Back/HyperHessians/test.jl

@@ -18,6 +18,7 @@ backends = [
 for backend in backends
     @test DI.check_available(backend)
     @test !DI.check_inplace(backend)
+    test_counterparts(backend)
 end
 
 scenarios = default_scenarios(; include_constantified = true, include_cachified = true)

DifferentiationInterface/test/Back/Mooncake/test.jl

@@ -19,6 +19,18 @@ backends = [
 for backend in backends
     @test check_available(backend)
     @test check_inplace(backend)
+    test_counterparts(backend)
+end
+
+@testset "Counterpart config" begin
+    config = Mooncake.Config(; friendly_tangents = true)
+    @test DifferentiationInterface.forward_counterpart(AutoMooncake()) === AutoMooncakeForward()
+    @test DifferentiationInterface.reverse_counterpart(AutoMooncakeForward()) === AutoMooncake()
+    # the Mooncake config must carry over to the counterpart
+    @test DifferentiationInterface.forward_counterpart(AutoMooncake(; config)) ===
+        AutoMooncakeForward(; config)
+    @test DifferentiationInterface.reverse_counterpart(AutoMooncakeForward(; config)) ===
+        AutoMooncake(; config)
 end
 
 test_differentiation(

DifferentiationInterface/test/Back/PolyesterForwardDiff/test.jl

@@ -19,6 +19,7 @@ backends = [
 for backend in backends
     @test check_available(backend)
     @test check_inplace(backend)
+    test_counterparts(backend)
     @test DifferentiationInterface.inner_preparation_behavior(backend) isa
         DifferentiationInterface.PrepareInnerOverload
 end

DifferentiationInterface/test/Back/ReverseDiff/test.jl

@@ -16,6 +16,7 @@ second_order_backends = [SecondOrder(AutoForwardDiff(), AutoReverseDiff())]
 for backend in vcat(backends, second_order_backends)
     @test check_available(backend)
     @test check_inplace(backend)
+    test_counterparts(backend)
 end
 
 ## Dense

DifferentiationInterface/test/Back/Symbolics/test.jl

@@ -12,6 +12,7 @@ check_no_implicit_imports(DifferentiationInterface)
 for backend in [AutoSymbolics(), AutoSparse(AutoSymbolics())]
     @test check_available(backend)
     @test check_inplace(backend)
+    test_counterparts(backend)
 end
 
 test_differentiation(

DifferentiationInterface/test/Back/Tracker/test.jl

@@ -10,6 +10,7 @@ check_no_implicit_imports(DifferentiationInterface)
 for backend in [AutoTracker()]
     @test check_available(backend)
     @test !check_inplace(backend)
+    test_counterparts(backend)
 end
 
 test_differentiation(

DifferentiationInterface/test/Back/Zygote/test.jl

@@ -17,6 +17,7 @@ second_order_backends = [SecondOrder(AutoForwardDiff(), AutoZygote())]
 for backend in vcat(backends, second_order_backends)
     @test check_available(backend)
     @test !check_inplace(backend)
+    test_counterparts(backend)
 end
 
 ## Dense