Use native gradient API for ForwardDiff, Enzyme, Mooncake by yebai · Pull Request #458 · TuringLang/Bijectors.jl

yebai · 2026-04-12T20:06:06Z

Sister PR for TuringLang/DynamicPPL.jl#1354

- Remove DifferentiationInterface from [deps]; add ADTypes - Move Enzyme to [weakdeps]; add BijectorsEnzymeExt extension - Add src/ad_utils.jl defining _value_and_gradient/_value_and_jacobian generic functions - Implement native backends in each pkg ext: ForwardDiff, ReverseDiff (compiled + non-compiled), Mooncake (reverse + forward JVP), Enzyme (reverse + forward) - Update src/vector/test_utils.jl to use ADTypes backend types and B._value_and_* API Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

github-actions · 2026-04-12T20:09:09Z

Bijectors.jl documentation for PR #458 is available at:
https://TuringLang.github.io/Bijectors.jl/previews/PR458/

- Avoid double f(x) evaluation in gradient/jacobian for ForwardDiff and ReverseDiff by using DiffResults (GradientResult, JacobianResult) with the in-place ! variants - For Enzyme reverse mode, use autodiff(ReverseWithPrimal, ...) to get value and gradient in one pass instead of calling f(x) separately - Fix _enzyme_mode to guard against mode=nothing (AutoEnzyme() default) which previously threw a MethodError from set_runtime_activity(::Nothing) - Pre-allocate dy/dx tangent buffers outside loops in Mooncake implementations and use fill! to zero them, avoiding one heap allocation per iteration - Add fallback _value_and_gradient/_value_and_jacobian methods with a clear error message for backends without a loaded extension Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

- Add AutoEnzyme{Nothing} to the ReverseWithPrimal dispatch union so the default (mode=nothing) backend also avoids double-evaluating f - Remove redundant `return` before `error(...)` in ad_utils.jl fallback methods; error() returns Union{} so return is a no-op Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

- Add return before error() in ad_utils.jl (JuliaFormatter) - Use ReverseDiff.DiffResults instead of ForwardDiff.DiffResults so the extension triggers on ReverseDiff alone - Keep Bijectors in test/Project.toml for B._value_and_jacobian calls Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

gdalle

There seems to be a lot of duplication of efforts compared to DI, along with some forgotten aspects. All in all, I'm not sure what we gain here

gdalle · 2026-04-15T10:09:11Z

+}
+
+function _annotate_function(f, backend::AutoEnzyme, mode)
+    annotation = typeof(backend).parameters[2]


Accessing type parameters this way is not recommended since the field is internal (AFAICT)

gdalle · 2026-04-15T10:10:13Z

+    EnzymeCore.Duplicated,EnzymeCore.DuplicatedNoNeed,EnzymeCore.MixedDuplicated
+}
+
+function _annotate_function(f, backend::AutoEnzyme, mode)


This looks a lot like https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/a5ecbe0b2bc97eaaac53a1c5c0b13c17f22f1ae9/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl#L42-L62, so I'm not sure where we save engineering effort

gdalle · 2026-04-15T10:10:41Z

+    backend::Union{AutoEnzyme{Nothing},AutoEnzyme{<:EnzymeCore.ReverseMode}},
+    x::AbstractVector,
+)
+    mode = if backend isa AutoEnzyme{Nothing}


See https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/a5ecbe0b2bc97eaaac53a1c5c0b13c17f22f1ae9/DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl#L161-L162

gdalle · 2026-04-15T10:11:02Z

+    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


Enzyme has a built-in forward-mode gradient function, which DI already uses in such cases. Any reason not to use it here too?
Ping @wsmoses

gdalle · 2026-04-15T10:11:48Z

+    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


Enzyme has a built-in forward Jacobian function, which DI already uses in such cases. Any reason not to use it here too?

gdalle · 2026-04-15T10:15:01Z

+    if T === Nothing
+        ForwardDiff.checktag(config, f, x)
+    end
+    ForwardDiff.gradient!(result, f, x, config, Val(false))


See https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/a5ecbe0b2bc97eaaac53a1c5c0b13c17f22f1ae9/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl#L395-L398

gdalle · 2026-04-15T10:15:51Z

+function _mooncake_zero_tangent_or_primal(
+    x, backend::Union{AutoMooncake,AutoMooncakeForward}
+)
+    if _mooncake_config(backend).friendly_tangents


This is type-unstable

gdalle · 2026-04-15T10:16:20Z

+        return f(x), similar(x, 0)
+    end
+    tape = ReverseDiff.GradientTape(f, x)
+    compiled = ReverseDiff.compile(tape)


Is it really worth compiling a tape you will only use once? I predict this slows things down significantly

gdalle · 2026-04-15T10:16:26Z

+
+function _value_and_jacobian(f, ::AutoReverseDiff{true}, x::AbstractVector)
+    tape = ReverseDiff.JacobianTape(f, x)
+    compiled = ReverseDiff.compile(tape)


gdalle · 2026-04-15T10:17:15Z

+
+Implementations are provided by package extensions for each AD backend.
+"""
+function _value_and_gradient(f, backend::ADTypes.AbstractADType, x::AbstractVector)


This is breaking

github-actions Bot assigned yebai Apr 12, 2026

yebai and others added 7 commits April 12, 2026 21:21

Apply JuliaFormatter v1.0.62

901c787

Fix native AD backend edge cases

5e3d21e

Use Mooncake value_and_gradient!! helpers

9ccf137

Remove DifferentiationInterface references

669d4a8

yebai force-pushed the native-ad-extensions branch from b3577f0 to 669d4a8 Compare April 12, 2026 23:15

yebai added 3 commits April 13, 2026 00:37

Clean extension imports and AD edge cases

a468093

Handle empty Enzyme reverse Jacobians

81bbaa4

Fix empty Enzyme Jacobian output shape

e4d8585

gdalle mentioned this pull request Apr 15, 2026

Tracking removal from Turing-verse JuliaDiff/DifferentiationInterface.jl#992

Closed

gdalle suggested changes Apr 15, 2026

View reviewed changes

yebai mentioned this pull request Apr 23, 2026

Use AbstractPPL AD interface in tests. #460

Open

Conversation

yebai commented Apr 12, 2026

Uh oh!

github-actions Bot commented Apr 12, 2026

Uh oh!

gdalle left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants