Use native gradient API for ForwardDiff, Enzyme, Mooncake

Project.toml

@@ -3,11 +3,11 @@ uuid = "76274a88-744f-5084-9051-94815aaf08c4"
 version = "0.15.20"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
@@ -28,29 +28,32 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 [weakdeps]
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 
 [extensions]
 BijectorsDistributionsADExt = "DistributionsAD"
+BijectorsEnzymeExt = "Enzyme"
 BijectorsForwardDiffExt = "ForwardDiff"
 BijectorsLazyArraysExt = "LazyArrays"
 BijectorsMooncakeExt = "Mooncake"
 BijectorsReverseDiffChainRulesExt = ["ChainRules", "ReverseDiff"]
 BijectorsReverseDiffExt = "ReverseDiff"
 
 [compat]
+ADTypes = "1"
 AbstractPPL = "0.14"
 ArgCheck = "1, 2"
 ChainRules = "1"
 ChainRulesCore = "0.10.11, 1"
 ChangesOfVariables = "0.1"
-DifferentiationInterface = "0.7.14"
 Distributions = "0.25.33"
 DistributionsAD = "0.6"
 DocStringExtensions = "0.9"
+Enzyme = "0.13"
 EnzymeCore = "0.8.15"
 FillArrays = "1"
 ForwardDiff = "0.10, 1.0.1"
@@ -60,7 +63,7 @@ IrrationalConstants = "0.1, 0.2"
 LazyArrays = "2"
 LogExpFunctions = "0.3.3"
 MappedArrays = "0.2.2, 0.3, 0.4"
-Mooncake = "0.4.95, 0.5"
+Mooncake = "0.5.26"
 Reexport = "0.2, 1"
 ReverseDiff = "1"
 Roots = "1.3.15, 2"
@@ -70,6 +73,7 @@ julia = "1.10.8"
 
 [extras]
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02"

docs/Project.toml

@@ -1,7 +1,6 @@
 [deps]
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"

docs/src/defining_examples.md

@@ -207,12 +207,10 @@ function full_transform(x12)
     return StereographicProj()(x123)
 end
 
-import DifferentiationInterface as DI
 using FiniteDifferences, LinearAlgebra
 x = [0.3, 0.4, sgn * sqrt(1 - 0.3^2 - 0.4^2)]
 
-adtype = DI.AutoFiniteDifferences(; fdm=central_fdm(5, 1))
-jac = DI.jacobian(full_transform, adtype, x[1:2])
+jac = only(FiniteDifferences.jacobian(central_fdm(5, 1), full_transform, x[1:2]))
 logjac = logabsdet(jac)[1]
 ```
 

ext/BijectorsEnzymeExt.jl

@@ -0,0 +1,105 @@
+module BijectorsEnzymeExt
+
+import Bijectors: _value_and_gradient, _value_and_jacobian
+import ADTypes: AutoEnzyme
+using Enzyme: Enzyme
+using EnzymeCore: EnzymeCore
+
+const DuplicatedFunctionAnnotations = Union{
+    EnzymeCore.Duplicated,EnzymeCore.DuplicatedNoNeed,EnzymeCore.MixedDuplicated
+}
+
+function _annotate_function(f, backend::AutoEnzyme, mode)
+    annotation = typeof(backend).parameters[2]
+    if annotation === Nothing
+        return f
+    elseif annotation <: EnzymeCore.Const
+        return Enzyme.Const(f)
+    elseif annotation <: DuplicatedFunctionAnnotations
+        if Enzyme.guess_activity(typeof(f), mode) <: EnzymeCore.Const
+            return Enzyme.Const(f)
+        else
+            # Enzyme's sugar APIs only preserve function shadows for `Duplicated`,
+            # so normalize the duplicated-like annotations here.
+            return Enzyme.Duplicated(f, Enzyme.make_zero(f))
+        end
+    else
+        throw(ArgumentError("unsupported Enzyme function annotation $annotation"))
+    end
+end
+
+function _value_and_gradient(
+    f,
+    backend::Union{AutoEnzyme{Nothing},AutoEnzyme{<:EnzymeCore.ReverseMode}},
+    x::AbstractVector,
+)
+    mode = if backend isa AutoEnzyme{Nothing}
+        Enzyme.ReverseWithPrimal
+    else
+        Enzyme.WithPrimal(backend.mode)
+    end
+    annotated_f = _annotate_function(f, backend, mode)
+    dx = zero(x)
+    _, val = Enzyme.autodiff(mode, annotated_f, Enzyme.Active, Enzyme.Duplicated(x, dx))
+    return val, dx
+end
+
+function _value_and_gradient(
+    f, backend::AutoEnzyme{<:EnzymeCore.ForwardMode}, x::AbstractVector
+)
+    mode = Enzyme.WithPrimal(backend.mode)
+    annotated_f = _annotate_function(f, backend, mode)
+    grad = zero(x)
+    value = f(x)
+    for i in eachindex(x)
+        dx = zero(x)
+        dx[i] = one(eltype(x))
+        directional, primal = Enzyme.autodiff(mode, annotated_f, Enzyme.Duplicated(x, dx))
+        grad[i] = directional
+        if i == firstindex(x)
+            value = primal
+        end
+    end
+    return value, grad
+end
+
+function _value_and_jacobian(
+    f, backend::AutoEnzyme{<:EnzymeCore.ReverseMode}, x::AbstractVector
+)
+    value = f(x)
+    if isempty(x) || isempty(value)
+        return value, Matrix{eltype(value)}(undef, length(value), length(x))
+    end
+    annotated_f = _annotate_function(f, backend, backend.mode)
+    jacobian = only(Enzyme.jacobian(backend.mode, annotated_f, x))
+    return value, reshape(jacobian, length(value), length(x))
+end
+
+function _value_and_jacobian(f, ::AutoEnzyme{Nothing}, x::AbstractVector)
+    return _value_and_jacobian(f, AutoEnzyme(; mode=Enzyme.Forward), x)
+end
+
+function _value_and_jacobian(
+    f, backend::AutoEnzyme{<:EnzymeCore.ForwardMode}, x::AbstractVector
+)
+    mode = Enzyme.WithPrimal(backend.mode)
+    annotated_f = _annotate_function(f, backend, mode)
+    value = f(x)
+    J = nothing
+    for i in eachindex(x)
+        dx = zero(x)
+        dx[i] = one(eltype(x))
+        directional, primal = Enzyme.autodiff(mode, annotated_f, Enzyme.Duplicated(x, dx))
+        if i == firstindex(x)
+            value = primal isa AbstractArray ? copy(primal) : primal
+            J = Matrix{eltype(directional)}(undef, length(directional), length(x))
+        end
+        J[:, i] .= directional
+    end
+    if isnothing(J)
+        J = Matrix{eltype(value)}(undef, length(value), 0)
+    end
+    return value, J
+end
+
+end

ext/BijectorsForwardDiffExt.jl

@@ -1,8 +1,44 @@
 module BijectorsForwardDiffExt
 
-using Bijectors: Bijectors, find_alpha
+import Bijectors: Bijectors, find_alpha, _value_and_gradient, _value_and_jacobian
+import ADTypes: AutoForwardDiff
 using ForwardDiff: ForwardDiff
 
+function _value_and_gradient(
+    f, backend::AutoForwardDiff{chunksize,T}, x::AbstractVector
+) where {chunksize,T}
+    if isempty(x)
+        return f(x), similar(x, 0)
+    end
+    result = ForwardDiff.DiffResults.GradientResult(x)
+    chunk = isnothing(chunksize) ? ForwardDiff.Chunk(x) : ForwardDiff.Chunk{chunksize}()
+    tag = T === Nothing ? ForwardDiff.Tag(f, eltype(x)) : backend.tag
+    config = ForwardDiff.GradientConfig(nothing, x, chunk, tag)
+    if T === Nothing
+        ForwardDiff.checktag(config, f, x)
+    end
+    ForwardDiff.gradient!(result, f, x, config, Val(false))
+    return ForwardDiff.DiffResults.value(result), ForwardDiff.DiffResults.gradient(result)
+end
+
+function _value_and_jacobian(
+    f, backend::AutoForwardDiff{chunksize,T}, x::AbstractVector
+) where {chunksize,T}
+    y = f(x)
+    if isempty(x)
+        return y, Matrix{eltype(y)}(undef, length(y), 0)
+    end
+    result = ForwardDiff.DiffResults.JacobianResult(y, x)
+    chunk = isnothing(chunksize) ? ForwardDiff.Chunk(x) : ForwardDiff.Chunk{chunksize}()
+    tag = T === Nothing ? ForwardDiff.Tag(f, eltype(x)) : backend.tag
+    config = ForwardDiff.JacobianConfig(nothing, x, chunk, tag)
+    if T === Nothing
+        ForwardDiff.checktag(config, f, x)
+    end
+    ForwardDiff.jacobian!(result, f, x, config, Val(false))
+    return ForwardDiff.DiffResults.value(result), ForwardDiff.DiffResults.jacobian(result)
+end
+
 Bijectors._eps(::Type{<:ForwardDiff.Dual{<:Any,Real}}) = Bijectors._eps(Real)
 Bijectors._eps(::Type{<:ForwardDiff.Dual{<:Any,<:Integer}}) = Bijectors._eps(Real)
 

ext/BijectorsLazyArraysExt.jl

@@ -1,15 +1,15 @@
 module BijectorsLazyArraysExt
 
 import Bijectors: maporbroadcast
-using LazyArrays: LazyArrays
+using LazyArrays: BroadcastArray
 
-function maporbroadcast(f, x1::LazyArrays.BroadcastArray, x...)
+function maporbroadcast(f, x1::BroadcastArray, x...)
     return copy(f.(x1, x...))
 end
-function maporbroadcast(f, x1, x2::LazyArrays.BroadcastArray, x...)
+function maporbroadcast(f, x1, x2::BroadcastArray, x...)
     return copy(f.(x1, x2, x...))
 end
-function maporbroadcast(f, x1, x2, x3::LazyArrays.BroadcastArray, x...)
+function maporbroadcast(f, x1, x2, x3::BroadcastArray, x...)
     return copy(f.(x1, x2, x3, x...))
 end
 

ext/BijectorsMooncakeExt.jl

@@ -1,8 +1,101 @@
 module BijectorsMooncakeExt
 
+using Mooncake: Mooncake
 using Mooncake:
-    @is_primitive, MinimalCtx, Mooncake, CoDual, primal, tangent_type, @from_chainrules
+    @is_primitive,
+    MinimalCtx,
+    CoDual,
+    tangent_type,
+    @from_chainrules,
+    prepare_pullback_cache,
+    prepare_gradient_cache,
+    prepare_derivative_cache,
+    value_and_pullback!!,
+    value_and_gradient!!,
+    value_and_derivative!!,
+    zero_tangent,
+    Config,
+    _copy_output,
+    tangent_to_primal!!
 using Bijectors: find_alpha, ChainRulesCore
+import Bijectors: _value_and_gradient, _value_and_jacobian
+import ADTypes: AutoMooncake, AutoMooncakeForward
+
+_mooncake_config(::Union{AutoMooncake{Nothing},AutoMooncakeForward{Nothing}}) = Config()
+_mooncake_config(backend::Union{AutoMooncake,AutoMooncakeForward}) = backend.config
+
+function _mooncake_zero_tangent_or_primal(
+    x, backend::Union{AutoMooncake,AutoMooncakeForward}
+)
+    if _mooncake_config(backend).friendly_tangents
+        return tangent_to_primal!!(_copy_output(x), zero_tangent(x))
+    else
+        return zero_tangent(x)
+    end
+end
+
+## Reverse-mode implementations
+
+function _value_and_gradient(f, backend::AutoMooncake, x::AbstractVector)
+    cache = prepare_gradient_cache(f, x; config=_mooncake_config(backend))
+    val, (_, x_grad) = value_and_gradient!!(cache, f, x)
+    return val, x_grad
+end
+
+function _value_and_jacobian(f, backend::AutoMooncake, x::AbstractVector)
+    cache = prepare_pullback_cache(f, x; config=_mooncake_config(backend))
+    n_out, n_in = length(cache.y_cache), length(x)
+    dy = zeros(eltype(cache.y_cache), n_out)
+    if n_out > 0
+        dy[1] = one(eltype(cache.y_cache))
+    end
+    val, (_, first_row) = value_and_pullback!!(cache, dy, f, x)
+    if n_out == 0
+        return _copy_output(val), Matrix{eltype(x)}(undef, 0, n_in)
+    end
+    y = _copy_output(val)
+    J = Matrix{eltype(first_row)}(undef, n_out, n_in)
+    J[1, :] .= first_row
+    for i in 2:n_out
+        fill!(dy, zero(eltype(cache.y_cache)))
+        dy[i] = one(eltype(cache.y_cache))
+        _, (_, row) = value_and_pullback!!(cache, dy, f, x)
+        J[i, :] .= row
+    end
+    return y, J
+end
+
+## Forward-mode implementations (column-by-column JVPs)
+
+function _value_and_gradient(f, backend::AutoMooncakeForward, x::AbstractVector)
+    cache = prepare_gradient_cache(f, x; config=_mooncake_config(backend))
+    val, (_, x_grad) = value_and_gradient!!(cache, f, x)
+    return val, x_grad
+end
+
+function _value_and_jacobian(
+    f, backend::AutoMooncakeForward, x::AbstractVector{T}
+) where {T}
+    y = f(x)
+    n_out, n_in = length(y), length(x)
+    cache = prepare_derivative_cache(f, x; config=_mooncake_config(backend))
+    df = _mooncake_zero_tangent_or_primal(f, backend)
+    if n_in == 0
+        return y, Matrix{eltype(y)}(undef, n_out, 0)
+    end
+    dx = zeros(T, n_in)
+    dx[1] = one(T)
+    _, first_jvp = value_and_derivative!!(cache, (f, df), (x, dx))
+    J = Matrix{eltype(first_jvp)}(undef, n_out, n_in)
+    J[:, 1] .= first_jvp
+    for j in 2:n_in
+        fill!(dx, zero(T))
+        dx[j] = one(T)
+        _, jvp = value_and_derivative!!(cache, (f, df), (x, dx))
+        J[:, j] .= jvp
+    end
+    return y, J
+end
 
 @from_chainrules(MinimalCtx, Tuple{typeof(find_alpha),Float16,Float16,Float16})
 @from_chainrules(MinimalCtx, Tuple{typeof(find_alpha),Float32,Float32,Float32})
@@ -50,7 +143,9 @@ function Mooncake.rrule!!(
         msg = "Integer argument has tangent type $(tangent_type(I)), should be NoTangent."
         throw(ArgumentError(msg))
     end
-    out, pb = ChainRulesCore.rrule(find_alpha, primal(x), primal(y), primal(z))
+    out, pb = ChainRulesCore.rrule(
+        find_alpha, Mooncake.primal(x), Mooncake.primal(y), Mooncake.primal(z)
+    )
     function find_alpha_pb(dout::P)
         _, dx, dy, _ = pb(dout)
         return Mooncake.NoRData(), P(dx), P(dy), Mooncake.NoRData()