diff --git a/DifferentiationInterface/CHANGELOG.md b/DifferentiationInterface/CHANGELOG.md index 93173bf13..df58c14b5 100644 --- a/DifferentiationInterface/CHANGELOG.md +++ b/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 diff --git a/DifferentiationInterface/docs/src/api.md b/DifferentiationInterface/docs/src/api.md index 108f1c03d..eba5cee96 100644 --- a/DifferentiationInterface/docs/src/api.md +++ b/DifferentiationInterface/docs/src/api.md @@ -123,6 +123,8 @@ DifferentiationInterface.inner ```@docs DifferentiateWith +DifferentiationInterface.forward_counterpart +DifferentiationInterface.reverse_counterpart ``` ### Sparsity tools diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl index dbf1e4f5c..5649a1d82 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/DifferentiationInterfaceEnzymeExt.jl +++ b/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") diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/counterparts.jl b/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/counterparts.jl new file mode 100644 index 000000000..4a28aad1e --- /dev/null +++ b/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 diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl index db3cdebbc..dcd912a9d 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl +++ b/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") diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/counterparts.jl b/DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/counterparts.jl new file mode 100644 index 000000000..2a1e6153b --- /dev/null +++ b/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) diff --git a/DifferentiationInterface/src/DifferentiationInterface.jl b/DifferentiationInterface/src/DifferentiationInterface.jl index 96b3ff0f6..6f60b65df 100644 --- a/DifferentiationInterface/src/DifferentiationInterface.jl +++ b/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") diff --git a/DifferentiationInterface/src/misc/from_primitive.jl b/DifferentiationInterface/src/misc/from_primitive.jl index 5bbd3490a..9caa84b31 100644 --- a/DifferentiationInterface/src/misc/from_primitive.jl +++ b/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 diff --git a/DifferentiationInterface/src/misc/zero_backends.jl b/DifferentiationInterface/src/misc/zero_backends.jl index c91d1b204..b2a4279ed 100644 --- a/DifferentiationInterface/src/misc/zero_backends.jl +++ b/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() diff --git a/DifferentiationInterface/src/utils/counterparts.jl b/DifferentiationInterface/src/utils/counterparts.jl new file mode 100644 index 000000000..cf832560e --- /dev/null +++ b/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. diff --git a/DifferentiationInterface/test/Back/ChainRules/test.jl b/DifferentiationInterface/test/Back/ChainRules/test.jl index b484a15cb..47e9212d8 100644 --- a/DifferentiationInterface/test/Back/ChainRules/test.jl +++ b/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( diff --git a/DifferentiationInterface/test/Back/Enzyme/test.jl b/DifferentiationInterface/test/Back/Enzyme/test.jl index d43019d43..0b087a654 100644 --- a/DifferentiationInterface/test/Back/Enzyme/test.jl +++ b/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 diff --git a/DifferentiationInterface/test/Back/FastDifferentiation/test.jl b/DifferentiationInterface/test/Back/FastDifferentiation/test.jl index e42dc5132..44078c3ee 100644 --- a/DifferentiationInterface/test/Back/FastDifferentiation/test.jl +++ b/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( diff --git a/DifferentiationInterface/test/Back/FiniteDiff/test.jl b/DifferentiationInterface/test/Back/FiniteDiff/test.jl index d2a63f03e..68201c3cc 100644 --- a/DifferentiationInterface/test/Back/FiniteDiff/test.jl +++ b/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 diff --git a/DifferentiationInterface/test/Back/FiniteDifferences/test.jl b/DifferentiationInterface/test/Back/FiniteDifferences/test.jl index ba585dc0c..26daa6555 100644 --- a/DifferentiationInterface/test/Back/FiniteDifferences/test.jl +++ b/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 diff --git a/DifferentiationInterface/test/Back/ForwardDiff/test.jl b/DifferentiationInterface/test/Back/ForwardDiff/test.jl index 59321b309..1b3fabaf2 100644 --- a/DifferentiationInterface/test/Back/ForwardDiff/test.jl +++ b/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 diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl index fb074f9f3..204a14a4a 100644 --- a/DifferentiationInterface/test/Back/GTPSA/test.jl +++ b/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) diff --git a/DifferentiationInterface/test/Back/HyperHessians/test.jl b/DifferentiationInterface/test/Back/HyperHessians/test.jl index 1378ca567..cf7ccb7d7 100644 --- a/DifferentiationInterface/test/Back/HyperHessians/test.jl +++ b/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) diff --git a/DifferentiationInterface/test/Back/Mooncake/test.jl b/DifferentiationInterface/test/Back/Mooncake/test.jl index eb80af551..3a90ba34a 100644 --- a/DifferentiationInterface/test/Back/Mooncake/test.jl +++ b/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( diff --git a/DifferentiationInterface/test/Back/PolyesterForwardDiff/test.jl b/DifferentiationInterface/test/Back/PolyesterForwardDiff/test.jl index 8ee877b98..67118465c 100644 --- a/DifferentiationInterface/test/Back/PolyesterForwardDiff/test.jl +++ b/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 diff --git a/DifferentiationInterface/test/Back/ReverseDiff/test.jl b/DifferentiationInterface/test/Back/ReverseDiff/test.jl index 2450b9bbf..c4021a381 100644 --- a/DifferentiationInterface/test/Back/ReverseDiff/test.jl +++ b/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 diff --git a/DifferentiationInterface/test/Back/Symbolics/test.jl b/DifferentiationInterface/test/Back/Symbolics/test.jl index 5bcf9830a..7d625eb8a 100644 --- a/DifferentiationInterface/test/Back/Symbolics/test.jl +++ b/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( diff --git a/DifferentiationInterface/test/Back/Tracker/test.jl b/DifferentiationInterface/test/Back/Tracker/test.jl index a6fcfb671..abf30adcf 100644 --- a/DifferentiationInterface/test/Back/Tracker/test.jl +++ b/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( diff --git a/DifferentiationInterface/test/Back/Zygote/test.jl b/DifferentiationInterface/test/Back/Zygote/test.jl index ceeeedbc5..4e570c95c 100644 --- a/DifferentiationInterface/test/Back/Zygote/test.jl +++ b/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 diff --git a/DifferentiationInterface/test/Core/Internals/backends.jl b/DifferentiationInterface/test/Core/Internals/backends.jl index 43bb4f1f3..f3d865351 100644 --- a/DifferentiationInterface/test/Core/Internals/backends.jl +++ b/DifferentiationInterface/test/Core/Internals/backends.jl @@ -1,8 +1,9 @@ using ADTypes -using ADTypes: mode +using ADTypes: dense_ad, mode using DifferentiationInterface using DifferentiationInterface: AutoSimpleFiniteDiff, + AutoForwardFromPrimitive, AutoReverseFromPrimitive, inner, outer, @@ -11,7 +12,9 @@ using DifferentiationInterface: inplace_support, pushforward_performance, pullback_performance, - hvp_mode + hvp_mode, + forward_counterpart, + reverse_counterpart import DifferentiationInterface as DI using Test @@ -21,6 +24,8 @@ rb = AutoReverseFromPrimitive(AutoSimpleFiniteDiff()) @testset "NoAutoDiff" begin @test_throws NoAutoDiffSelectedError check_available(NoAutoDiff()) @test_throws NoAutoDiffSelectedError inplace_support(NoAutoDiff()) + @test_throws NoAutoDiffSelectedError forward_counterpart(NoAutoDiff()) + @test_throws NoAutoDiffSelectedError reverse_counterpart(NoAutoDiff()) end @testset "SecondOrder" begin @@ -32,6 +37,8 @@ end @test Bool(inplace_support(backend)) @test_throws ArgumentError pushforward_performance(backend) @test_throws ArgumentError pullback_performance(backend) + @test_throws ArgumentError forward_counterpart(backend) + @test_throws ArgumentError reverse_counterpart(backend) end @testset "MixedMode" begin @@ -43,6 +50,21 @@ end @test Bool(inplace_support(backend)) @test_throws MethodError pushforward_performance(backend) @test_throws MethodError pullback_performance(backend) + # a MixedMode backend is its own counterpart in both directions + @test forward_counterpart(backend) === backend + @test reverse_counterpart(backend) === backend +end + +@testset "Counterparts" begin + # forward-/reverse-mode backends are their own counterpart + @test forward_counterpart(fb) === fb + @test reverse_counterpart(rb) === rb + # without a known counterpart, the backend is returned unchanged, with a warning + @test (@test_logs (:warn, r"reverse-mode counterpart") reverse_counterpart(fb)) === fb + # FromPrimitive wrappers swap the primitive, applying the counterpart inside + @test forward_counterpart(rb) isa AutoForwardFromPrimitive{<:Any, <:AutoSimpleFiniteDiff} + @test reverse_counterpart(forward_counterpart(rb)) isa + AutoReverseFromPrimitive{<:Any, <:AutoSimpleFiniteDiff} end @testset "Sparse" begin @@ -54,4 +76,11 @@ end @test_throws ArgumentError pullback_performance(backend) @test_throws ArgumentError hvp_mode(backend) end + # counterparts act on the dense backend and preserve the sparsity machinery + backend = AutoSparse(rb) + counterpart = forward_counterpart(backend) + @test counterpart isa AutoSparse + @test dense_ad(counterpart) isa AutoForwardFromPrimitive + @test counterpart.sparsity_detector === backend.sparsity_detector + @test counterpart.coloring_algorithm === backend.coloring_algorithm end diff --git a/DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl b/DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl index 40f460320..d67e11577 100644 --- a/DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl +++ b/DifferentiationInterfaceTest/src/DifferentiationInterfaceTest.jl @@ -87,7 +87,10 @@ using DifferentiationInterface: outer, inplace_support, pushforward_performance, - pullback_performance + pullback_performance, + check_available, + forward_counterpart, + reverse_counterpart using DifferentiationInterface: Rewrap, Context, Constant, Cache, ConstantOrCache, unwrap using DifferentiationInterface: PreparationMismatchError using DocStringExtensions: TYPEDFIELDS, TYPEDSIGNATURES @@ -137,6 +140,7 @@ include("scenarios/empty.jl") include("scenarios/extensions.jl") include("tests/correctness_eval.jl") +include("tests/counterparts.jl") include("tests/prep_eval.jl") include("tests/type_stability.jl") include("tests/benchmark.jl") @@ -147,6 +151,7 @@ include("test_differentiation.jl") export FIRST_ORDER, SECOND_ORDER export Scenario, compute_results export test_differentiation, benchmark_differentiation +export test_counterparts export DifferentiationBenchmarkDataRow @compile_workload begin diff --git a/DifferentiationInterfaceTest/src/tests/counterparts.jl b/DifferentiationInterfaceTest/src/tests/counterparts.jl new file mode 100644 index 000000000..8e65d0949 --- /dev/null +++ b/DifferentiationInterfaceTest/src/tests/counterparts.jl @@ -0,0 +1,31 @@ +const FORWARD_CAPABLE_MODES = Union{ForwardMode, ForwardOrReverseMode, SymbolicMode} +const REVERSE_CAPABLE_MODES = Union{ReverseMode, ForwardOrReverseMode, SymbolicMode} + +""" + test_counterparts(backend) + +Test that `forward_counterpart` and `reverse_counterpart` behave sensibly for `backend`: +the returned backends are available, they support the requested mode (unless `backend` is +returned unchanged for lack of a known counterpart), and applying the same counterpart +again leaves the mode unchanged. +""" +function test_counterparts(backend::AbstractADType) + @testset "Counterparts: $(typeof(backend))" begin + fc = forward_counterpart(backend) + rc = reverse_counterpart(backend) + @test check_available(fc) + @test check_available(rc) + # the counterpart must support the requested mode, except when the backend was + # returned unchanged for lack of a known counterpart + if fc !== backend || mode(backend) isa FORWARD_CAPABLE_MODES + @test mode(fc) isa FORWARD_CAPABLE_MODES + end + if rc !== backend || mode(backend) isa REVERSE_CAPABLE_MODES + @test mode(rc) isa REVERSE_CAPABLE_MODES + end + # applying the counterpart a second time leaves the mode unchanged + @test mode(forward_counterpart(fc)) == mode(fc) + @test mode(reverse_counterpart(rc)) == mode(rc) + end + return nothing +end diff --git a/DifferentiationInterfaceTest/test/zero_backends.jl b/DifferentiationInterfaceTest/test/zero_backends.jl index 0bdafa3a2..7b84c5783 100644 --- a/DifferentiationInterfaceTest/test/zero_backends.jl +++ b/DifferentiationInterfaceTest/test/zero_backends.jl @@ -7,6 +7,11 @@ using DifferentiationInterfaceTest using DifferentiationInterfaceTest: allocfree_scenarios, no_matrices using Test +## Counterparts + +test_counterparts(AutoZeroForward()) +test_counterparts(AutoZeroReverse()) + ## Type stability test_differentiation(