Support AutoEnzyme Jacobians inside Reactant kernels

.github/workflows/Test.yml

@@ -95,6 +95,7 @@ jobs:
           - DifferentiateWith
           # - Diffractor
           - Enzyme
+          - EnzymeReactant
           - FastDifferentiation
           - FiniteDiff
           - FiniteDifferences

DifferentiationInterface/Project.toml

@@ -23,6 +23,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
@@ -49,6 +50,7 @@ DifferentiationInterfacePolyesterForwardDiffExt = [
     "ForwardDiff",
     "DiffResults",
 ]
+DifferentiationInterfaceReactantExt = "Reactant"
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
 DifferentiationInterfaceSparseArraysExt = "SparseArrays"
 DifferentiationInterfaceSparseConnectivityTracerExt = "SparseConnectivityTracer"
@@ -76,6 +78,7 @@ HyperHessians = "0.3"
 LinearAlgebra = "1"
 Mooncake = "0.5.25"
 PolyesterForwardDiff = "0.1.2"
+Reactant = "0.2.283"
 ReverseDiff = "1.15.1"
 SparseArrays = "1"
 SparseConnectivityTracer = "0.6.14, 1"

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -216,6 +216,40 @@ end
 
 ## Jacobian
 
+struct EnzymeReactantJacobianPrep{SIG, D, P} <: DI.JacobianPrep{SIG}
+    _sig::Val{SIG}
+    directions::D
+    pushforward_prep::P
+end
+
+function prepare_reactant_jacobian(
+        strict::Val,
+        f_or_f!y::Tuple,
+        backend::AutoEnzyme,
+        x,
+        contexts::Vararg{DI.Context, C}
+    ) where {C}
+    _sig = DI.signature(f_or_f!y..., backend, x, contexts...; strict)
+    directions = onehot(x)
+    pushforward_prep = DI.prepare_pushforward_nokwarg(
+        strict, f_or_f!y..., backend, x, directions, contexts...
+    )
+    return EnzymeReactantJacobianPrep(_sig, directions, pushforward_prep)
+end
+
+function reactant_jacobian(
+        f_or_f!y::Tuple,
+        prep::EnzymeReactantJacobianPrep,
+        backend::AutoEnzyme,
+        x,
+        contexts::Vararg{DI.Context, C}
+    ) where {C}
+    columns = DI.pushforward(
+        f_or_f!y..., prep.pushforward_prep, backend, x, prep.directions, contexts...
+    )
+    return DI.stack_vec_col(columns)
+end
+
 struct EnzymeForwardOneArgJacobianPrep{SIG, B, DF, DC, O} <: DI.JacobianPrep{SIG}
     _sig::Val{SIG}
     _valB::Val{B}
@@ -232,6 +266,9 @@ function DI.prepare_jacobian_nokwarg(
         x,
         contexts::Vararg{DI.Constant, C}
     ) where {F, C}
+    if DI._use_reactant_jacobian(backend)
+        return prepare_reactant_jacobian(strict, (f,), backend, x, contexts...)
+    end
     _sig = DI.signature(f, backend, x, contexts...; strict)
     y = f(x, map(DI.unwrap, contexts)...)
     valB = to_val(DI.pick_batchsize(backend, x))
@@ -244,6 +281,29 @@ function DI.prepare_jacobian_nokwarg(
     )
 end
 
+function DI.jacobian(
+        f::F,
+        prep::EnzymeReactantJacobianPrep,
+        backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},
+        x,
+        contexts::Vararg{DI.Context, C},
+    ) where {F, C}
+    DI.check_prep(f, prep, backend, x, contexts...)
+    return reactant_jacobian((f,), prep, backend, x, contexts...)
+end
+
+function DI.jacobian!(
+        f::F,
+        jac,
+        prep::EnzymeReactantJacobianPrep,
+        backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},
+        x,
+        contexts::Vararg{DI.Context, C},
+    ) where {F, C}
+    DI.check_prep(f, prep, backend, x, contexts...)
+    return copyto!(jac, reactant_jacobian((f,), prep, backend, x, contexts...))
+end
+
 function DI.jacobian(
         f::F,
         prep::EnzymeForwardOneArgJacobianPrep{SIG, B},

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_twoarg.jl

@@ -114,3 +114,97 @@ function DI.pushforward!(
     DI.value_and_pushforward!(f!, y, ty, prep, backend, x, tx, contexts...)
     return ty
 end
+
+## Jacobian
+
+struct EnzymeReactantTwoArgJacobianPrep{SIG, D, P} <: DI.JacobianPrep{SIG}
+    _sig::Val{SIG}
+    directions::D
+    pushforward_prep::P
+end
+
+function reactant_out_of_place(x, f!::F, y, contexts...) where {F}
+    new_y = zero(y)
+    f!(new_y, x, contexts...)
+    return new_y
+end
+
+function prepare_reactant_twoarg_jacobian(
+        strict::Val,
+        f!::F,
+        y,
+        backend::AutoEnzyme,
+        x,
+        contexts::Vararg{DI.Context, C}
+    ) where {F, C}
+    _sig = DI.signature(f!, y, backend, x, contexts...; strict)
+    directions = onehot(x)
+    wrapped_contexts = (DI.Constant(f!), DI.Constant(y), contexts...)
+    pushforward_prep = DI.prepare_pushforward_nokwarg(
+        strict, reactant_out_of_place, backend, x, directions, wrapped_contexts...
+    )
+    return EnzymeReactantTwoArgJacobianPrep(_sig, directions, pushforward_prep)
+end
+
+function reactant_twoarg_jacobian(
+        f!::F,
+        y,
+        prep::EnzymeReactantTwoArgJacobianPrep,
+        backend::AutoEnzyme,
+        x,
+        contexts::Vararg{DI.Context, C}
+    ) where {F, C}
+    wrapped_contexts = (DI.Constant(f!), DI.Constant(y), contexts...)
+    columns = DI.pushforward(
+        reactant_out_of_place,
+        prep.pushforward_prep,
+        backend,
+        x,
+        prep.directions,
+        wrapped_contexts...,
+    )
+    return DI.stack_vec_col(columns)
+end
+
+function DI.prepare_jacobian_nokwarg(
+        strict::Val,
+        f!::F,
+        y,
+        backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},
+        x,
+        contexts::Vararg{DI.Context, C}
+    ) where {F, C}
+    if DI._use_reactant_jacobian(backend)
+        return prepare_reactant_twoarg_jacobian(strict, f!, y, backend, x, contexts...)
+    end
+    batch_size_settings = DI.pick_batchsize(backend, x)
+    return DI._prepare_jacobian_aux(
+        strict, DI.PushforwardFast(), batch_size_settings, y, (f!, y), backend, x, contexts...
+    )
+end
+
+function DI.jacobian(
+        f!::F,
+        y,
+        prep::EnzymeReactantTwoArgJacobianPrep,
+        backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},
+        x,
+        contexts::Vararg{DI.Context, C},
+    ) where {F, C}
+    DI.check_prep(f!, y, prep, backend, x, contexts...)
+    return reactant_twoarg_jacobian(f!, y, prep, backend, x, contexts...)
+end
+
+function DI.jacobian!(
+        f!::F,
+        y,
+        jac,
+        prep::EnzymeReactantTwoArgJacobianPrep,
+        backend::AutoEnzyme{<:Union{ForwardMode, Nothing}},
+        x,
+        contexts::Vararg{DI.Context, C},
+    ) where {F, C}
+    DI.check_prep(f!, y, prep, backend, x, contexts...)
+    new_jac = reactant_twoarg_jacobian(f!, y, prep, backend, x, contexts...)
+    return copyto!(jac, new_jac)
+end

DifferentiationInterface/ext/DifferentiationInterfaceReactantExt.jl

@@ -0,0 +1,9 @@
+module DifferentiationInterfaceReactantExt
+
+using ADTypes: AutoEnzyme
+import DifferentiationInterface as DI
+using Reactant: within_compile
+
+DI._use_reactant_jacobian(::AutoEnzyme) = within_compile()
+
+end

DifferentiationInterface/src/first_order/jacobian.jl

@@ -131,6 +131,8 @@ end
 
 ## Preparation
 
+_use_reactant_jacobian(::AbstractADType) = false
+
 abstract type StandardJacobianPrep{SIG} <: JacobianPrep{SIG} end
 
 struct PushforwardJacobianPrep{

DifferentiationInterface/test/Back/EnzymeReactant/Project.toml

@@ -0,0 +1,9 @@
+[deps]
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
+DifferentiationInterfaceTest = "a82114a7-5aa3-49a8-9643-716bb13727a3"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+DifferentiationInterface = { path = "../../.." }

DifferentiationInterface/test/Back/EnzymeReactant/test.jl

@@ -0,0 +1,71 @@
+using DifferentiationInterface:
+    AutoEnzyme,
+    Constant,
+    jacobian,
+    jacobian!,
+    prepare_jacobian,
+    value_and_jacobian
+using Enzyme: Enzyme
+using Reactant: Reactant, @jit
+using Test
+
+f(x, p) = x .^ 2 .- p
+f!(y, x, p) = y .= f(x, p)
+
+function oop_jacobian(x, p)
+    backend = AutoEnzyme()
+    prep = prepare_jacobian(f, backend, x, Constant(p))
+    return jacobian(f, prep, backend, x, Constant(p))
+end
+
+function oop_value_and_jacobian(x, p)
+    backend = AutoEnzyme()
+    prep = prepare_jacobian(f, backend, x, Constant(p))
+    return value_and_jacobian(f, prep, backend, x, Constant(p))
+end
+
+function iip_jacobian(x, p)
+    backend = AutoEnzyme()
+    y = zero(x)
+    prep = prepare_jacobian(f!, y, backend, x, Constant(p))
+    return jacobian(f!, y, prep, backend, x, Constant(p))
+end
+
+function iip_value_and_jacobian(x, p)
+    backend = AutoEnzyme()
+    y = zero(x)
+    prep = prepare_jacobian(f!, y, backend, x, Constant(p))
+    return value_and_jacobian(f!, y, prep, backend, x, Constant(p))
+end
+
+function oop_jacobian!(x, p)
+    backend = AutoEnzyme()
+    jac = similar(x, length(x), length(x))
+    prep = prepare_jacobian(f, backend, x, Constant(p))
+    return jacobian!(f, jac, prep, backend, x, Constant(p))
+end
+
+function iip_jacobian!(x, p)
+    backend = AutoEnzyme()
+    y = zero(x)
+    jac = similar(x, length(x), length(x))
+    prep = prepare_jacobian(f!, y, backend, x, Constant(p))
+    return jacobian!(f!, y, jac, prep, backend, x, Constant(p))
+end
+
+@testset "AutoEnzyme Jacobian inside Reactant" begin
+    x = Reactant.to_rarray(Float32[1, 2])
+    p = Reactant.to_rarray(Float32[3, 4])
+    expected_jacobian = Float32[2 0; 0 4]
+
+    @test @jit(oop_jacobian(x, p)) ≈ expected_jacobian
+    value, jac = @jit oop_value_and_jacobian(x, p)
+    @test value ≈ Float32[-2, 0]
+    @test jac ≈ expected_jacobian
+    @test @jit(iip_jacobian(x, p)) ≈ expected_jacobian
+    value, jac = @jit iip_value_and_jacobian(x, p)
+    @test value ≈ Float32[-2, 0]
+    @test jac ≈ expected_jacobian
+    @test @jit(oop_jacobian!(x, p)) ≈ expected_jacobian
+    @test @jit(iip_jacobian!(x, p)) ≈ expected_jacobian
+end