From 44d57ab23175491ee7e226aabb43ba73ec66541c Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Sun, 12 Apr 2026 21:04:57 +0100
Subject: [PATCH 01/11] Replace DifferentiationInterface with native AD package
 extensions

- 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>
---
 Project.toml                   |  8 +++-
 ext/BijectorsEnzymeExt.jl      | 34 +++++++++++++++++
 ext/BijectorsForwardDiffExt.jl | 13 ++++++-
 ext/BijectorsMooncakeExt.jl    | 69 +++++++++++++++++++++++++++++++++-
 ext/BijectorsReverseDiffExt.jl | 30 +++++++++++++++
 src/Bijectors.jl               |  2 +
 src/ad_utils.jl                | 20 ++++++++++
 src/vector/test_utils.jl       | 60 +++++++++++++++++------------
 8 files changed, 207 insertions(+), 29 deletions(-)
 create mode 100644 ext/BijectorsEnzymeExt.jl
 create mode 100644 src/ad_utils.jl

diff --git a/Project.toml b/Project.toml
index ae0beacb8..6cee9a60c 100644
--- a/Project.toml
+++ b/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,6 +28,7 @@ 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"
@@ -35,6 +36,7 @@ ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 
 [extensions]
 BijectorsDistributionsADExt = "DistributionsAD"
+BijectorsEnzymeExt = "Enzyme"
 BijectorsForwardDiffExt = "ForwardDiff"
 BijectorsLazyArraysExt = "LazyArrays"
 BijectorsMooncakeExt = "Mooncake"
@@ -42,15 +44,16 @@ 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"
@@ -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"
diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl
new file mode 100644
index 000000000..880fbbe8e
--- /dev/null
+++ b/ext/BijectorsEnzymeExt.jl
@@ -0,0 +1,34 @@
+module BijectorsEnzymeExt
+
+import Bijectors: _value_and_gradient, _value_and_jacobian
+import ADTypes: AutoEnzyme
+using Enzyme: Enzyme
+using EnzymeCore: EnzymeCore
+
+function _enzyme_mode(backend::AutoEnzyme)
+    return Enzyme.set_runtime_activity(backend.mode)
+end
+
+function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}) where {M,A}
+    if A <: EnzymeCore.Const
+        return Enzyme.Const(f)
+    else
+        return f
+    end
+end
+
+function _value_and_gradient(f, backend::AutoEnzyme, x::AbstractVector)
+    mode = _enzyme_mode(backend)
+    annotated_f = _enzyme_annotate_f(f, backend)
+    grads = Enzyme.gradient(mode, annotated_f, x)
+    return f(x), first(grads)
+end
+
+function _value_and_jacobian(f, backend::AutoEnzyme, x::AbstractVector)
+    mode = _enzyme_mode(backend)
+    annotated_f = _enzyme_annotate_f(f, backend)
+    jacs = Enzyme.jacobian(mode, annotated_f, x)
+    return f(x), first(jacs)
+end
+
+end
diff --git a/ext/BijectorsForwardDiffExt.jl b/ext/BijectorsForwardDiffExt.jl
index 76ae0dd2f..a7732ebe0 100644
--- a/ext/BijectorsForwardDiffExt.jl
+++ b/ext/BijectorsForwardDiffExt.jl
@@ -1,8 +1,19 @@
 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, ::AutoForwardDiff, x::AbstractVector)
+    grad = ForwardDiff.gradient(f, x)
+    return f(x), grad
+end
+
+function _value_and_jacobian(f, ::AutoForwardDiff, x::AbstractVector)
+    jac = ForwardDiff.jacobian(f, x)
+    return f(x), jac
+end
+
 Bijectors._eps(::Type{<:ForwardDiff.Dual{<:Any,Real}}) = Bijectors._eps(Real)
 Bijectors._eps(::Type{<:ForwardDiff.Dual{<:Any,<:Integer}}) = Bijectors._eps(Real)
 
diff --git a/ext/BijectorsMooncakeExt.jl b/ext/BijectorsMooncakeExt.jl
index bea5c2433..484529bed 100644
--- a/ext/BijectorsMooncakeExt.jl
+++ b/ext/BijectorsMooncakeExt.jl
@@ -1,8 +1,75 @@
 module BijectorsMooncakeExt
 
 using Mooncake:
-    @is_primitive, MinimalCtx, Mooncake, CoDual, primal, tangent_type, @from_chainrules
+    @is_primitive,
+    MinimalCtx,
+    Mooncake,
+    CoDual,
+    primal,
+    tangent_type,
+    @from_chainrules,
+    prepare_pullback_cache,
+    prepare_derivative_cache,
+    value_and_pullback!!,
+    value_and_derivative!!,
+    zero_tangent
 using Bijectors: find_alpha, ChainRulesCore
+import Bijectors: _value_and_gradient, _value_and_jacobian
+import ADTypes: AutoMooncake, AutoMooncakeForward
+
+## Reverse-mode implementations
+
+function _value_and_gradient(f, ::AutoMooncake, x::AbstractVector{T}) where {T}
+    cache = prepare_pullback_cache(f, x)
+    val, (_, x_grad) = value_and_pullback!!(cache, one(T), f, x)
+    return val, x_grad
+end
+
+function _value_and_jacobian(f, ::AutoMooncake, x::AbstractVector{T}) where {T}
+    y = f(x)
+    n_out, n_in = length(y), length(x)
+    J = Matrix{T}(undef, n_out, n_in)
+    cache = prepare_pullback_cache(f, x)
+    for i in 1:n_out
+        dy = zeros(eltype(y), n_out)
+        dy[i] = one(eltype(y))
+        _, (_, 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, ::AutoMooncakeForward, x::AbstractVector{T}) where {T}
+    val = f(x)
+    n = length(x)
+    grad = Vector{T}(undef, n)
+    cache = prepare_derivative_cache(f, x)
+    df = zero_tangent(f)
+    for j in 1:n
+        dx = zeros(T, n)
+        dx[j] = one(T)
+        _, jvp = value_and_derivative!!(cache, (f, df), (x, dx))
+        grad[j] = jvp
+    end
+    return val, grad
+end
+
+function _value_and_jacobian(f, ::AutoMooncakeForward, x::AbstractVector{T}) where {T}
+    y = f(x)
+    n_out, n_in = length(y), length(x)
+    J = Matrix{T}(undef, n_out, n_in)
+    cache = prepare_derivative_cache(f, x)
+    df = zero_tangent(f)
+    for j in 1:n_in
+        dx = zeros(T, n_in)
+        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})
diff --git a/ext/BijectorsReverseDiffExt.jl b/ext/BijectorsReverseDiffExt.jl
index ea5b959e2..99b8ede6b 100644
--- a/ext/BijectorsReverseDiffExt.jl
+++ b/ext/BijectorsReverseDiffExt.jl
@@ -10,6 +10,36 @@ using ReverseDiff:
     TrackedMatrix,
     @grad_from_chainrules
 
+import Bijectors: _value_and_gradient, _value_and_jacobian
+import ADTypes: AutoReverseDiff
+
+function _value_and_gradient(f, ::AutoReverseDiff{false}, x::AbstractVector)
+    grad = ReverseDiff.gradient(f, x)
+    return f(x), grad
+end
+
+function _value_and_jacobian(f, ::AutoReverseDiff{false}, x::AbstractVector)
+    jac = ReverseDiff.jacobian(f, x)
+    return f(x), jac
+end
+
+function _value_and_gradient(f, ::AutoReverseDiff{true}, x::AbstractVector)
+    tape = ReverseDiff.GradientTape(f, x)
+    compiled = ReverseDiff.compile(tape)
+    result = similar(x)
+    ReverseDiff.gradient!(result, compiled, x)
+    return f(x), result
+end
+
+function _value_and_jacobian(f, ::AutoReverseDiff{true}, x::AbstractVector)
+    tape = ReverseDiff.JacobianTape(f, x)
+    compiled = ReverseDiff.compile(tape)
+    y = f(x)
+    result = zeros(eltype(x), length(y), length(x))
+    ReverseDiff.jacobian!(result, compiled, x)
+    return y, result
+end
+
 using Bijectors:
     ChainRulesCore,
     Elementwise,
diff --git a/src/Bijectors.jl b/src/Bijectors.jl
index 3fd0ddba5..5c50692db 100644
--- a/src/Bijectors.jl
+++ b/src/Bijectors.jl
@@ -40,6 +40,7 @@ using InverseFunctions: InverseFunctions
 import ChangesOfVariables: ChangesOfVariables, with_logabsdet_jacobian
 import InverseFunctions: inverse
 
+using ADTypes: ADTypes
 using ChainRulesCore: ChainRulesCore
 using Functors: Functors
 using IrrationalConstants: IrrationalConstants
@@ -359,6 +360,7 @@ end
 include("utils.jl")
 include("interface.jl")
 include("chainrules.jl")
+include("ad_utils.jl")
 include("vector/VectorBijectors.jl")
 
 # Broadcasting here breaks some AD backends for certain array types
diff --git a/src/ad_utils.jl b/src/ad_utils.jl
new file mode 100644
index 000000000..7a9a46643
--- /dev/null
+++ b/src/ad_utils.jl
@@ -0,0 +1,20 @@
+"""
+    _value_and_gradient(f, backend::ADTypes.AbstractADType, x::AbstractVector)
+
+Compute the value and gradient of scalar-valued `f` at `x`.
+Returns `(f(x), gradient)` where `gradient` is a vector of the same size as `x`.
+
+Implementations are provided by package extensions for each AD backend.
+"""
+function _value_and_gradient end
+
+"""
+    _value_and_jacobian(f, backend::ADTypes.AbstractADType, x::AbstractVector)
+
+Compute the value and Jacobian of `f` at `x`.
+Returns `(f(x), jacobian)` where `jacobian` is a matrix of size
+`(length(f(x)), length(x))`.
+
+Implementations are provided by package extensions for each AD backend.
+"""
+function _value_and_jacobian end
diff --git a/src/vector/test_utils.jl b/src/vector/test_utils.jl
index bb5003f49..273f9d976 100644
--- a/src/vector/test_utils.jl
+++ b/src/vector/test_utils.jl
@@ -1,19 +1,25 @@
 using Test
 using LinearAlgebra: logabsdet, Cholesky, UpperTriangular, LowerTriangular
-import DifferentiationInterface as DI
+import ADTypes:
+    AutoForwardDiff,
+    AutoReverseDiff,
+    AutoMooncake,
+    AutoMooncakeForward,
+    AutoEnzyme,
+    AbstractADType
 import EnzymeCore as EC
 
 # Would like to use FiniteDifferences, but very easy to run into issues with
 # https://juliadiff.org/FiniteDifferences.jl/latest/#Dealing-with-Singularities
-const ref_adtype = DI.AutoForwardDiff()
-
-const default_adtypes = [
-    DI.AutoReverseDiff(),
-    DI.AutoReverseDiff(; compile=true),
-    DI.AutoMooncake(),
-    DI.AutoMooncakeForward(),
-    DI.AutoEnzyme(; mode=EC.Forward, function_annotation=EC.Const),
-    DI.AutoEnzyme(; mode=EC.Reverse, function_annotation=EC.Const),
+const ref_adtype = AutoForwardDiff()
+
+const default_adtypes = AbstractADType[
+    AutoReverseDiff(),
+    AutoReverseDiff(; compile=true),
+    AutoMooncake(; config=nothing),
+    AutoMooncakeForward(; config=nothing),
+    AutoEnzyme(; mode=EC.Forward, function_annotation=EC.Const),
+    AutoEnzyme(; mode=EC.Reverse, function_annotation=EC.Const),
 ]
 
 _get_value_support(::D.Distribution{<:Any,VS}) where {VS<:D.ValueSupport} = VS
@@ -422,7 +428,7 @@ function test_optics(d::D.Distribution)
         x = rand(d)
         xvec = to_vec(d)(x)
         yvec = to_linked_vec(d)(x)
-        J = DI.jacobian(to_linked_vec(d) ∘ from_vec(d), ref_adtype, xvec)
+        _, J = B._value_and_jacobian(to_linked_vec(d) ∘ from_vec(d), ref_adtype, xvec)
         o = optic_vec(d)
         lo = linked_optic_vec(d)
         for i in 1:length(yvec)
@@ -566,7 +572,9 @@ function test_logjac(d::D.Distribution, atol, rtol)
                 # to make sure that the Jacobian is square.
                 ad_xvec = to_vec_for_logjac_test(d)(x)
                 ad_ffwd = to_linked_vec(d) ∘ from_vec_for_logjac_test(d)
-                ad_logjac = first(logabsdet(DI.jacobian(ad_ffwd, ref_adtype, ad_xvec)))
+                ad_logjac = first(
+                    logabsdet(last(B._value_and_jacobian(ad_ffwd, ref_adtype, ad_xvec)))
+                )
                 @test vbt_logjac ≈ ad_logjac atol = atol rtol = rtol
             end
 
@@ -579,7 +587,9 @@ function test_logjac(d::D.Distribution, atol, rtol)
                 # For the AD calculation we need to use to/from_vec_for_logjac_test instead,
                 # to make sure that the Jacobian is square.
                 ad_frvs = to_vec_for_logjac_test(d) ∘ from_linked_vec(d)
-                ad_logjac = first(logabsdet(DI.jacobian(ad_frvs, ref_adtype, yvec)))
+                ad_logjac = first(
+                    logabsdet(last(B._value_and_jacobian(ad_frvs, ref_adtype, yvec)))
+                )
                 @test vbt_logjac ≈ ad_logjac atol = atol rtol = rtol
             end
         end
@@ -590,7 +600,7 @@ end
 Test that various AD backends can differentiate the conversions to and from vector and
 linked vector forms for the given distribution `d`.
 """
-function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, rtol)
+function test_ad(d::D.Distribution, adtypes::Vector{<:AbstractADType}, atol, rtol)
     # If `d` is a discrete distribution, Mooncake refuses to differentiate through the
     # transforms (which are just identity transforms). Likewise, Enzyme will throw an
     # error saying that the output is Const but was not marked as such.
@@ -600,9 +610,9 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol,
     adtypes = if d isa D.Distribution{<:Any,D.Discrete}
         filter(adtypes) do adtype
             !(
-                adtype isa DI.AutoMooncake ||
-                adtype isa DI.AutoMooncakeForward ||
-                adtype isa DI.AutoEnzyme
+                adtype isa AutoMooncake ||
+                adtype isa AutoMooncakeForward ||
+                adtype isa AutoEnzyme
             )
         end
     else
@@ -613,18 +623,18 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol,
         x = _rand_safe_ad(d)
         xvec = to_vec(d)(x)
         ffwd = to_linked_vec(d) ∘ from_vec(d)
-        ref_jac = DI.jacobian(ffwd, ref_adtype, xvec)
+        _, ref_jac = B._value_and_jacobian(ffwd, ref_adtype, xvec)
 
         ladj(xvec) = last(with_logabsdet_jacobian(ffwd, xvec))
-        ref_grad_ladj = DI.gradient(ladj, ref_adtype, xvec)
+        _, ref_grad_ladj = B._value_and_gradient(ladj, ref_adtype, xvec)
 
         for adtype in adtypes
             @testset let x = x, adtype = adtype, d = d
-                ad_jac = DI.jacobian(ffwd, adtype, xvec)
+                _, ad_jac = B._value_and_jacobian(ffwd, adtype, xvec)
                 @test ref_jac ≈ ad_jac atol = atol rtol = rtol
             end
             @testset let x = x, adtype = adtype, d = d
-                ad_grad_ladj = DI.gradient(ladj, adtype, xvec)
+                _, ad_grad_ladj = B._value_and_gradient(ladj, adtype, xvec)
                 @test ref_grad_ladj ≈ ad_grad_ladj atol = atol rtol = rtol
             end
         end
@@ -634,19 +644,19 @@ function test_ad(d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol,
         x = _rand_safe_ad(d)
         yvec = to_linked_vec(d)(x)
         frvs = to_vec(d) ∘ from_linked_vec(d)
-        ref_jac = DI.jacobian(frvs, ref_adtype, yvec)
+        _, ref_jac = B._value_and_jacobian(frvs, ref_adtype, yvec)
 
         ladj(yvec) = last(with_logabsdet_jacobian(frvs, yvec))
-        ref_grad_ladj = DI.gradient(ladj, ref_adtype, yvec)
+        _, ref_grad_ladj = B._value_and_gradient(ladj, ref_adtype, yvec)
 
         for adtype in adtypes
             @testset let x = x, adtype = adtype, d = d
-                ad_jac = DI.jacobian(frvs, adtype, yvec)
+                _, ad_jac = B._value_and_jacobian(frvs, adtype, yvec)
                 @test ref_jac ≈ ad_jac atol = atol rtol = rtol
             end
 
             @testset let x = x, adtype = adtype, d = d
-                ad_grad_ladj = DI.gradient(ladj, adtype, yvec)
+                _, ad_grad_ladj = B._value_and_gradient(ladj, adtype, yvec)
                 @test ref_grad_ladj ≈ ad_grad_ladj atol = atol rtol = rtol
             end
         end

From 16888407264383e25c59d79165fc576b2f71e0f4 Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Sun, 12 Apr 2026 21:21:47 +0100
Subject: [PATCH 02/11] Fix review issues in native AD extensions

- 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>
---
 ext/BijectorsEnzymeExt.jl      | 23 +++++++++++++++++++++++
 ext/BijectorsForwardDiffExt.jl | 11 +++++++----
 ext/BijectorsMooncakeExt.jl    |  9 ++++++---
 ext/BijectorsReverseDiffExt.jl | 19 +++++++++++--------
 src/ad_utils.jl                | 14 ++++++++++++--
 5 files changed, 59 insertions(+), 17 deletions(-)

diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl
index 880fbbe8e..e8173a5d1 100644
--- a/ext/BijectorsEnzymeExt.jl
+++ b/ext/BijectorsEnzymeExt.jl
@@ -5,10 +5,15 @@ import ADTypes: AutoEnzyme
 using Enzyme: Enzyme
 using EnzymeCore: EnzymeCore
 
+# When mode is unspecified, fall back to Reverse.
+function _enzyme_mode(::AutoEnzyme{Nothing})
+    return Enzyme.set_runtime_activity(Enzyme.Reverse)
+end
 function _enzyme_mode(backend::AutoEnzyme)
     return Enzyme.set_runtime_activity(backend.mode)
 end
 
+# Returns f annotated as requested, or bare f when annotation is Nothing.
 function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}) where {M,A}
     if A <: EnzymeCore.Const
         return Enzyme.Const(f)
@@ -17,6 +22,24 @@ function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}) where {M,A}
     end
 end
 
+# For explicit reverse mode, use ReverseWithPrimal so value and gradient are computed
+# in a single autodiff call rather than evaluating f a second time.
+function _value_and_gradient(
+    f, backend::AutoEnzyme{<:EnzymeCore.ReverseMode}, x::AbstractVector
+)
+    annotated_f = _enzyme_annotate_f(f, backend)
+    dx = zero(x)
+    _, val = Enzyme.autodiff(
+        Enzyme.set_runtime_activity(Enzyme.ReverseWithPrimal),
+        annotated_f,
+        Enzyme.Active,
+        Enzyme.Duplicated(x, dx),
+    )
+    return val, dx
+end
+
+# For forward mode (or auto), the gradient already requires O(n) JVPs; one extra
+# f(x) evaluation is negligible.
 function _value_and_gradient(f, backend::AutoEnzyme, x::AbstractVector)
     mode = _enzyme_mode(backend)
     annotated_f = _enzyme_annotate_f(f, backend)
diff --git a/ext/BijectorsForwardDiffExt.jl b/ext/BijectorsForwardDiffExt.jl
index a7732ebe0..60075b7d9 100644
--- a/ext/BijectorsForwardDiffExt.jl
+++ b/ext/BijectorsForwardDiffExt.jl
@@ -5,13 +5,16 @@ import ADTypes: AutoForwardDiff
 using ForwardDiff: ForwardDiff
 
 function _value_and_gradient(f, ::AutoForwardDiff, x::AbstractVector)
-    grad = ForwardDiff.gradient(f, x)
-    return f(x), grad
+    result = ForwardDiff.DiffResults.GradientResult(x)
+    ForwardDiff.gradient!(result, f, x)
+    return ForwardDiff.DiffResults.value(result), ForwardDiff.DiffResults.gradient(result)
 end
 
 function _value_and_jacobian(f, ::AutoForwardDiff, x::AbstractVector)
-    jac = ForwardDiff.jacobian(f, x)
-    return f(x), jac
+    y = f(x)
+    result = ForwardDiff.DiffResults.JacobianResult(y, x)
+    ForwardDiff.jacobian!(result, f, x)
+    return ForwardDiff.DiffResults.value(result), ForwardDiff.DiffResults.jacobian(result)
 end
 
 Bijectors._eps(::Type{<:ForwardDiff.Dual{<:Any,Real}}) = Bijectors._eps(Real)
diff --git a/ext/BijectorsMooncakeExt.jl b/ext/BijectorsMooncakeExt.jl
index 484529bed..0225e2b64 100644
--- a/ext/BijectorsMooncakeExt.jl
+++ b/ext/BijectorsMooncakeExt.jl
@@ -30,8 +30,9 @@ function _value_and_jacobian(f, ::AutoMooncake, x::AbstractVector{T}) where {T}
     n_out, n_in = length(y), length(x)
     J = Matrix{T}(undef, n_out, n_in)
     cache = prepare_pullback_cache(f, x)
+    dy = zeros(eltype(y), n_out)
     for i in 1:n_out
-        dy = zeros(eltype(y), n_out)
+        fill!(dy, zero(eltype(y)))
         dy[i] = one(eltype(y))
         _, (_, row) = value_and_pullback!!(cache, dy, f, x)
         J[i, :] .= row
@@ -47,8 +48,9 @@ function _value_and_gradient(f, ::AutoMooncakeForward, x::AbstractVector{T}) whe
     grad = Vector{T}(undef, n)
     cache = prepare_derivative_cache(f, x)
     df = zero_tangent(f)
+    dx = zeros(T, n)
     for j in 1:n
-        dx = zeros(T, n)
+        fill!(dx, zero(T))
         dx[j] = one(T)
         _, jvp = value_and_derivative!!(cache, (f, df), (x, dx))
         grad[j] = jvp
@@ -62,8 +64,9 @@ function _value_and_jacobian(f, ::AutoMooncakeForward, x::AbstractVector{T}) whe
     J = Matrix{T}(undef, n_out, n_in)
     cache = prepare_derivative_cache(f, x)
     df = zero_tangent(f)
+    dx = zeros(T, n_in)
     for j in 1:n_in
-        dx = zeros(T, n_in)
+        fill!(dx, zero(T))
         dx[j] = one(T)
         _, jvp = value_and_derivative!!(cache, (f, df), (x, dx))
         J[:, j] .= jvp
diff --git a/ext/BijectorsReverseDiffExt.jl b/ext/BijectorsReverseDiffExt.jl
index 99b8ede6b..4f56e3dd4 100644
--- a/ext/BijectorsReverseDiffExt.jl
+++ b/ext/BijectorsReverseDiffExt.jl
@@ -14,30 +14,33 @@ import Bijectors: _value_and_gradient, _value_and_jacobian
 import ADTypes: AutoReverseDiff
 
 function _value_and_gradient(f, ::AutoReverseDiff{false}, x::AbstractVector)
-    grad = ReverseDiff.gradient(f, x)
-    return f(x), grad
+    result = ReverseDiff.DiffResults.GradientResult(x)
+    ReverseDiff.gradient!(result, f, x)
+    return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.gradient(result)
 end
 
 function _value_and_jacobian(f, ::AutoReverseDiff{false}, x::AbstractVector)
-    jac = ReverseDiff.jacobian(f, x)
-    return f(x), jac
+    y = f(x)
+    result = ReverseDiff.DiffResults.JacobianResult(y, x)
+    ReverseDiff.jacobian!(result, f, x)
+    return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.jacobian(result)
 end
 
 function _value_and_gradient(f, ::AutoReverseDiff{true}, x::AbstractVector)
     tape = ReverseDiff.GradientTape(f, x)
     compiled = ReverseDiff.compile(tape)
-    result = similar(x)
+    result = ReverseDiff.DiffResults.GradientResult(x)
     ReverseDiff.gradient!(result, compiled, x)
-    return f(x), result
+    return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.gradient(result)
 end
 
 function _value_and_jacobian(f, ::AutoReverseDiff{true}, x::AbstractVector)
     tape = ReverseDiff.JacobianTape(f, x)
     compiled = ReverseDiff.compile(tape)
     y = f(x)
-    result = zeros(eltype(x), length(y), length(x))
+    result = ReverseDiff.DiffResults.JacobianResult(y, x)
     ReverseDiff.jacobian!(result, compiled, x)
-    return y, result
+    return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.jacobian(result)
 end
 
 using Bijectors:
diff --git a/src/ad_utils.jl b/src/ad_utils.jl
index 7a9a46643..a4982005c 100644
--- a/src/ad_utils.jl
+++ b/src/ad_utils.jl
@@ -6,7 +6,12 @@ Returns `(f(x), gradient)` where `gradient` is a vector of the same size as `x`.
 
 Implementations are provided by package extensions for each AD backend.
 """
-function _value_and_gradient end
+function _value_and_gradient(f, backend::ADTypes.AbstractADType, x::AbstractVector)
+    return error(
+        "_value_and_gradient is not implemented for backend $(typeof(backend)). " *
+        "Load the corresponding AD package to enable support.",
+    )
+end
 
 """
     _value_and_jacobian(f, backend::ADTypes.AbstractADType, x::AbstractVector)
@@ -17,4 +22,9 @@ Returns `(f(x), jacobian)` where `jacobian` is a matrix of size
 
 Implementations are provided by package extensions for each AD backend.
 """
-function _value_and_jacobian end
+function _value_and_jacobian(f, backend::ADTypes.AbstractADType, x::AbstractVector)
+    return error(
+        "_value_and_jacobian is not implemented for backend $(typeof(backend)). " *
+        "Load the corresponding AD package to enable support.",
+    )
+end

From 0b6b6d3becddd1758a310e9a8699c9de39b9268e Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Sun, 12 Apr 2026 21:26:21 +0100
Subject: [PATCH 03/11] Fix Enzyme Nothing dispatch and remove return from
 error calls

- 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>
---
 ext/BijectorsEnzymeExt.jl | 12 +++++++-----
 src/ad_utils.jl           |  4 ++--
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl
index e8173a5d1..ccf880f46 100644
--- a/ext/BijectorsEnzymeExt.jl
+++ b/ext/BijectorsEnzymeExt.jl
@@ -22,10 +22,12 @@ function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}) where {M,A}
     end
 end
 
-# For explicit reverse mode, use ReverseWithPrimal so value and gradient are computed
-# in a single autodiff call rather than evaluating f a second time.
+# For reverse mode (explicit or auto/nothing), use ReverseWithPrimal so value and
+# gradient are computed in a single autodiff call rather than evaluating f twice.
 function _value_and_gradient(
-    f, backend::AutoEnzyme{<:EnzymeCore.ReverseMode}, x::AbstractVector
+    f,
+    backend::Union{AutoEnzyme{Nothing},AutoEnzyme{<:EnzymeCore.ReverseMode}},
+    x::AbstractVector,
 )
     annotated_f = _enzyme_annotate_f(f, backend)
     dx = zero(x)
@@ -38,8 +40,8 @@ function _value_and_gradient(
     return val, dx
 end
 
-# For forward mode (or auto), the gradient already requires O(n) JVPs; one extra
-# f(x) evaluation is negligible.
+# For forward mode the gradient already requires O(n) JVPs; one extra f(x) evaluation
+# is negligible.
 function _value_and_gradient(f, backend::AutoEnzyme, x::AbstractVector)
     mode = _enzyme_mode(backend)
     annotated_f = _enzyme_annotate_f(f, backend)
diff --git a/src/ad_utils.jl b/src/ad_utils.jl
index a4982005c..b4207673e 100644
--- a/src/ad_utils.jl
+++ b/src/ad_utils.jl
@@ -7,7 +7,7 @@ Returns `(f(x), gradient)` where `gradient` is a vector of the same size as `x`.
 Implementations are provided by package extensions for each AD backend.
 """
 function _value_and_gradient(f, backend::ADTypes.AbstractADType, x::AbstractVector)
-    return error(
+    error(
         "_value_and_gradient is not implemented for backend $(typeof(backend)). " *
         "Load the corresponding AD package to enable support.",
     )
@@ -23,7 +23,7 @@ Returns `(f(x), jacobian)` where `jacobian` is a matrix of size
 Implementations are provided by package extensions for each AD backend.
 """
 function _value_and_jacobian(f, backend::ADTypes.AbstractADType, x::AbstractVector)
-    return error(
+    error(
         "_value_and_jacobian is not implemented for backend $(typeof(backend)). " *
         "Load the corresponding AD package to enable support.",
     )

From 9774644c460d77e7ed232db9e406cd1d44a64fbe Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Sun, 12 Apr 2026 21:33:11 +0100
Subject: [PATCH 04/11] Fix formatter and revert ReverseDiff ext to use
 ReverseDiff.DiffResults

- 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>
---
 docs/src/defining.md                        |  8 ++++----
 docs/src/defining_examples.md               |  6 +++---
 docs/src/vector.md                          |  3 ++-
 src/ad_utils.jl                             |  4 ++--
 src/bijectors/named_bijector.jl             |  4 ++--
 src/bijectors/named_stacked.jl              | 13 +++++++------
 src/vector/product/product.jl               |  7 +++----
 src/vector/reshaped/reshaped.jl             |  6 ++----
 test/Project.toml                           |  1 +
 test/ad/enzyme.jl                           |  9 +++------
 test/bijectors/product_bijector.jl          |  4 ++--
 test/bijectors/rational_quadratic_spline.jl |  2 +-
 12 files changed, 32 insertions(+), 35 deletions(-)

diff --git a/docs/src/defining.md b/docs/src/defining.md
index a92c81bdf..67b94a4ba 100644
--- a/docs/src/defining.md
+++ b/docs/src/defining.md
@@ -8,14 +8,14 @@ In general, there are two pieces of information needed to define a bijector:
 
  2. The log-absolute determinant of the Jacobian of that transformation.
     For a transformation $b: \mathbb{R}^d \to \mathbb{R}^d$, the Jacobian at point $x \in \mathbb{R}^d$ is defined as:
-    
+
     $$J_{b}(x) = \begin{bmatrix}
     \partial y_1/\partial x_1 & \partial y_1/\partial x_2 & \cdots & \partial y_1/\partial x_d \\
     \partial y_2/\partial x_1 & \partial y_2/\partial x_2 & \cdots & \partial y_2/\partial x_d \\
     \vdots & \vdots & \ddots & \vdots \\
     \partial y_d/\partial x_1 & \partial y_d/\partial x_2 & \cdots & \partial y_d/\partial x_d
     \end{bmatrix}$$
-    
+
     where $y = b(x)$.
 
 ## The transform itself
@@ -27,7 +27,7 @@ Bijectors.with_logabsdet_jacobian
 ```
 
 !!! note
-    
+
     `with_logabsdet_jacobian` is re-exported from ChangesOfVariables.jl, so if you want to avoid importing Bijectors.jl, you can implement `ChangesOfVariables.with_logabsdet_jacobian` instead.
 
 If you define `with_logabsdet_jacobian(b, x)`, then you will automatically get default implementations of both `transform(b, x)` and `logabsdetjac(b, x)`, which respectively return the first and second value of that tuple.
@@ -60,7 +60,7 @@ Bijectors.inverse
 ```
 
 !!! note
-    
+
     `inverse` is re-exported from InverseFunctions.jl, so the same note as for `with_logabsdet_jacobian` applies.
 
 If `b` is a bijector, then `inverse(b)` should return the inverse bijector $b^{-1}$.
diff --git a/docs/src/defining_examples.md b/docs/src/defining_examples.md
index 95556b844..ac30efb43 100644
--- a/docs/src/defining_examples.md
+++ b/docs/src/defining_examples.md
@@ -94,7 +94,7 @@ inverse(b)(y) == x
 ```
 
 !!! note
-    
+
     Bijectors re-exports both `with_logabsdet_jacobian` as well as `inverse`, so you don't need to import them separately if Bijectors is already a dependency.
     Conversely, if you don't want to depend on Bijectors.jl directly, you can just import these functions from their respective packages.
 
@@ -121,7 +121,7 @@ end
 ```
 
 !!! warning
-    
+
     This will return `[Inf, Inf]` if `x[3] == 1` (the 'north pole' of the sphere), which may potentially make downstream computations fail. One potential way around this is to add `eps(T)` to the denominator to avoid it ever being zero: you will sometimes see this trick used in Bijectors.jl. However, be aware that the reverse transform has to also be modified accordingly so that the two transforms remain inverses of each other!
 
 When it comes to computing the Jacobian, we find ourselves in a spot of bother.
@@ -171,7 +171,7 @@ $$J = \begin{bmatrix}
 \end{bmatrix}.$$
 
 !!! note
-    
+
     When you see $x_3$ here, don't think 'the variable $x_3$': it's just shorthand for $\pm \sqrt{1 - x_1^2 - x_2^2}$. (And recall that these formulae hold for both choices of sign.)
 
 Its determinant very nicely simplifies to
diff --git a/docs/src/vector.md b/docs/src/vector.md
index 655d7dd4b..83f90d908 100644
--- a/docs/src/vector.md
+++ b/docs/src/vector.md
@@ -7,8 +7,9 @@ It assumes that there are three forms of samples from a distribution `d` that we
  1. **The original form**, which is what `rand(d)` returns.
 
  2. **A vectorised form**, which is a vector that contains a flattened version of the original form.
+
  3. **A linked vectorised form**, which is a vector in which:
-    
+
       + each element is independent; and
       + each element is unconstrained (can take any value in ℝ).
 
diff --git a/src/ad_utils.jl b/src/ad_utils.jl
index b4207673e..a4982005c 100644
--- a/src/ad_utils.jl
+++ b/src/ad_utils.jl
@@ -7,7 +7,7 @@ Returns `(f(x), gradient)` where `gradient` is a vector of the same size as `x`.
 Implementations are provided by package extensions for each AD backend.
 """
 function _value_and_gradient(f, backend::ADTypes.AbstractADType, x::AbstractVector)
-    error(
+    return error(
         "_value_and_gradient is not implemented for backend $(typeof(backend)). " *
         "Load the corresponding AD package to enable support.",
     )
@@ -23,7 +23,7 @@ Returns `(f(x), jacobian)` where `jacobian` is a matrix of size
 Implementations are provided by package extensions for each AD backend.
 """
 function _value_and_jacobian(f, backend::ADTypes.AbstractADType, x::AbstractVector)
-    error(
+    return error(
         "_value_and_jacobian is not implemented for backend $(typeof(backend)). " *
         "Load the corresponding AD package to enable support.",
     )
diff --git a/src/bijectors/named_bijector.jl b/src/bijectors/named_bijector.jl
index 85e53e9bc..4838a177d 100644
--- a/src/bijectors/named_bijector.jl
+++ b/src/bijectors/named_bijector.jl
@@ -139,7 +139,7 @@ deps(b::NamedCoupling{<:Any,Deps}) where {Deps} = Deps
     return quote
         b = nc.f($([:(x.$d) for d in deps]...))
         x_target, logjac = with_logabsdet_jacobian(b, x.$target)
-        return merge(x, ($target=x_target,)), logjac
+        return merge(x, (($target)=x_target,)), logjac
     end
 end
 
@@ -149,6 +149,6 @@ end
     return quote
         ib = inverse(ni.orig.f($([:(x.$d) for d in deps]...)))
         x_target, logjac = with_logabsdet_jacobian(ib, x.$target)
-        return merge(x, ($target=x_target,)), logjac
+        return merge(x, (($target)=x_target,)), logjac
     end
 end
diff --git a/src/bijectors/named_stacked.jl b/src/bijectors/named_stacked.jl
index 59076ac88..67d790b77 100644
--- a/src/bijectors/named_stacked.jl
+++ b/src/bijectors/named_stacked.jl
@@ -87,8 +87,8 @@ end
                 end
             ),
         )
-        push!(exprs, :(transforms = merge(transforms, ($n=trf,))))
-        push!(exprs, :(ranges = merge(ranges, ($n=output_range,))))
+        push!(exprs, :(transforms = merge(transforms, (($n)=trf,))))
+        push!(exprs, :(ranges = merge(ranges, (($n)=output_range,))))
     end
     push!(exprs, :(return NamedStacked{names}(transforms, ranges)))
     return Expr(:block, exprs...)
@@ -153,14 +153,15 @@ end
         if i == 1
             push!(
                 exprs,
-                :(output = ($n=inverse(nsi.orig.transforms.$n)(y[nsi.orig.ranges.$n]),)),
+                :(output = (($n)=inverse(nsi.orig.transforms.$n)(y[nsi.orig.ranges.$n]),)),
             )
         else
             push!(
                 exprs,
                 :(
                     output = merge(
-                        output, ($n=inverse(nsi.orig.transforms.$n)(y[nsi.orig.ranges.$n]),)
+                        output,
+                        (($n)=inverse(nsi.orig.transforms.$n)(y[nsi.orig.ranges.$n]),),
                     )
                 ),
             )
@@ -182,7 +183,7 @@ end
                     first_out, first_logjac = with_logabsdet_jacobian(
                         inverse(nsi.orig.transforms.$n), y[nsi.orig.ranges.$n]
                     )
-                    output = ($n=first_out,)
+                    output = (($n)=first_out,)
                     logjac = first_logjac
                 end,
             )
@@ -193,7 +194,7 @@ end
                     next_out, next_logjac = with_logabsdet_jacobian(
                         inverse(nsi.orig.transforms.$n), y[nsi.orig.ranges.$n]
                     )
-                    output = merge(output, ($n=next_out,))
+                    output = merge(output, (($n)=next_out,))
                     logjac += next_logjac
                 end,
             )
diff --git a/src/vector/product/product.jl b/src/vector/product/product.jl
index 33c784a3b..1539355fb 100644
--- a/src/vector/product/product.jl
+++ b/src/vector/product/product.jl
@@ -54,8 +54,7 @@ end
 for struct_type in (:ProductVecTransform, :ProductVecInvTransform)
     @eval begin
         Base.:(==)(t1::$struct_type, t2::$struct_type) =
-            (t1.transforms == t2.transforms) &
-            (t1.ranges == t2.ranges) &
+            (t1.transforms == t2.transforms) & (t1.ranges == t2.ranges) &
             (t1.base_size == t2.base_size)
         Base.isequal(t1::$struct_type, t2::$struct_type) =
             isequal(t1.transforms, t2.transforms) &&
@@ -303,7 +302,7 @@ end
             expr.args,
             :((xr, lj) = with_logabsdet_jacobian(t.transforms.$nm, view(y, t.ranges.$nm))),
         )
-        push!(expr.args, :(x = (x..., $nm=xr)))
+        push!(expr.args, :(x = (x..., ($nm)=xr)))
         push!(expr.args, :(logjac += lj))
     end
     push!(expr.args, :(return x, logjac))
@@ -332,7 +331,7 @@ end
     push!(exprs, :(offset = 1))
     for nm in names
         push!(exprs, :(this_length = length_fn(dists.$nm)))
-        push!(exprs, :(ranges = (ranges..., $nm=offset:(offset + this_length - 1))))
+        push!(exprs, :(ranges = (ranges..., ($nm)=offset:(offset + this_length - 1))))
         push!(exprs, :(offset += this_length))
     end
     push!(exprs, :(return struct_type(trfms, ranges, size(dists[1]))))
diff --git a/src/vector/reshaped/reshaped.jl b/src/vector/reshaped/reshaped.jl
index a829406e2..8cb8e1840 100644
--- a/src/vector/reshaped/reshaped.jl
+++ b/src/vector/reshaped/reshaped.jl
@@ -28,8 +28,7 @@ struct ReshapeWrapper{N1,N2,T1<:NTuple{N1,Int},T2<:NTuple{N2,Int},B}
     bijector::B
 end
 function Base.:(==)(r1::ReshapeWrapper, r2::ReshapeWrapper)
-    return (r1.reshaped_size == r2.reshaped_size) &
-           (r1.original_size == r2.original_size) &
+    return (r1.reshaped_size == r2.reshaped_size) & (r1.original_size == r2.original_size) &
            (r1.bijector == r2.bijector)
 end
 function Base.isequal(r1::ReshapeWrapper, r2::ReshapeWrapper)
@@ -62,8 +61,7 @@ struct InvReshapeWrapper{N1,N2,T1<:NTuple{N1,Int},T2<:NTuple{N2,Int},B}
     inv_bijector::B
 end
 function Base.:(==)(r1::InvReshapeWrapper, r2::InvReshapeWrapper)
-    return (r1.reshaped_size == r2.reshaped_size) &
-           (r1.original_size == r2.original_size) &
+    return (r1.reshaped_size == r2.reshaped_size) & (r1.original_size == r2.original_size) &
            (r1.inv_bijector == r2.inv_bijector)
 end
 function Base.isequal(r1::InvReshapeWrapper, r2::InvReshapeWrapper)
diff --git a/test/Project.toml b/test/Project.toml
index 8f6b34696..aeb288b9b 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -2,6 +2,7 @@
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
+Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
diff --git a/test/ad/enzyme.jl b/test/ad/enzyme.jl
index 39ef5b177..29a6e8983 100644
--- a/test/ad/enzyme.jl
+++ b/test/ad/enzyme.jl
@@ -32,8 +32,7 @@ using Test
         @testset "forward" begin
             # No batches
             @testset for RT in (Const, Duplicated, DuplicatedNoNeed),
-                Tx in (Const, Duplicated),
-                Ty in (Const, Duplicated),
+                Tx in (Const, Duplicated), Ty in (Const, Duplicated),
                 Tz in (Const, Duplicated)
 
                 test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
@@ -41,8 +40,7 @@ using Test
 
             # Batches
             @testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed),
-                Tx in (Const, BatchDuplicated),
-                Ty in (Const, BatchDuplicated),
+                Tx in (Const, BatchDuplicated), Ty in (Const, BatchDuplicated),
                 Tz in (Const, BatchDuplicated)
 
                 test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
@@ -51,8 +49,7 @@ using Test
         @testset "reverse" begin
             # No batches
             @testset for RT in (Const, Active),
-                Tx in (Const, Active),
-                Ty in (Const, Active),
+                Tx in (Const, Active), Ty in (Const, Active),
                 Tz in (Const, Active)
 
                 test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
diff --git a/test/bijectors/product_bijector.jl b/test/bijectors/product_bijector.jl
index 818f89c09..73dbba086 100644
--- a/test/bijectors/product_bijector.jl
+++ b/test/bijectors/product_bijector.jl
@@ -39,7 +39,7 @@ has_square_jacobian(b, x) = Bijectors.output_size(b, x) == size(x)
             y,
             logjac,
             changes_of_variables_test=has_square_jacobian(b, xs[1]),
-            test_not_identity=!isidentity,
+            test_not_identity=(!isidentity),
         )
     end
 
@@ -63,7 +63,7 @@ has_square_jacobian(b, x) = Bijectors.output_size(b, x) == size(x)
             y,
             logjac,
             changes_of_variables_test=has_square_jacobian(b, xs[1]),
-            test_not_identity=!isidentity,
+            test_not_identity=(!isidentity),
         )
     end
 end
diff --git a/test/bijectors/rational_quadratic_spline.jl b/test/bijectors/rational_quadratic_spline.jl
index a936c01c8..cfb315d81 100644
--- a/test/bijectors/rational_quadratic_spline.jl
+++ b/test/bijectors/rational_quadratic_spline.jl
@@ -46,7 +46,7 @@ using LogExpFunctions
 
         # Outside of domain
         x = 5.0
-        test_bijector(b, -x; y=-x, logjac=0)
+        test_bijector(b, -x; y=(-x), logjac=0)
         test_bijector(b, x; y=x, logjac=0)
 
         # multivariate

From 901c78706e164326a4f7c1c83177291ac72e714d Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Sun, 12 Apr 2026 22:45:40 +0100
Subject: [PATCH 05/11] Apply JuliaFormatter v1.0.62

---
 docs/src/defining.md            | 8 ++++----
 docs/src/defining_examples.md   | 6 +++---
 docs/src/vector.md              | 3 +--
 src/vector/product/product.jl   | 3 ++-
 src/vector/reshaped/reshaped.jl | 6 ++++--
 test/ad/enzyme.jl               | 9 ++++++---
 6 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/docs/src/defining.md b/docs/src/defining.md
index 67b94a4ba..a92c81bdf 100644
--- a/docs/src/defining.md
+++ b/docs/src/defining.md
@@ -8,14 +8,14 @@ In general, there are two pieces of information needed to define a bijector:
 
  2. The log-absolute determinant of the Jacobian of that transformation.
     For a transformation $b: \mathbb{R}^d \to \mathbb{R}^d$, the Jacobian at point $x \in \mathbb{R}^d$ is defined as:
-
+    
     $$J_{b}(x) = \begin{bmatrix}
     \partial y_1/\partial x_1 & \partial y_1/\partial x_2 & \cdots & \partial y_1/\partial x_d \\
     \partial y_2/\partial x_1 & \partial y_2/\partial x_2 & \cdots & \partial y_2/\partial x_d \\
     \vdots & \vdots & \ddots & \vdots \\
     \partial y_d/\partial x_1 & \partial y_d/\partial x_2 & \cdots & \partial y_d/\partial x_d
     \end{bmatrix}$$
-
+    
     where $y = b(x)$.
 
 ## The transform itself
@@ -27,7 +27,7 @@ Bijectors.with_logabsdet_jacobian
 ```
 
 !!! note
-
+    
     `with_logabsdet_jacobian` is re-exported from ChangesOfVariables.jl, so if you want to avoid importing Bijectors.jl, you can implement `ChangesOfVariables.with_logabsdet_jacobian` instead.
 
 If you define `with_logabsdet_jacobian(b, x)`, then you will automatically get default implementations of both `transform(b, x)` and `logabsdetjac(b, x)`, which respectively return the first and second value of that tuple.
@@ -60,7 +60,7 @@ Bijectors.inverse
 ```
 
 !!! note
-
+    
     `inverse` is re-exported from InverseFunctions.jl, so the same note as for `with_logabsdet_jacobian` applies.
 
 If `b` is a bijector, then `inverse(b)` should return the inverse bijector $b^{-1}$.
diff --git a/docs/src/defining_examples.md b/docs/src/defining_examples.md
index ac30efb43..95556b844 100644
--- a/docs/src/defining_examples.md
+++ b/docs/src/defining_examples.md
@@ -94,7 +94,7 @@ inverse(b)(y) == x
 ```
 
 !!! note
-
+    
     Bijectors re-exports both `with_logabsdet_jacobian` as well as `inverse`, so you don't need to import them separately if Bijectors is already a dependency.
     Conversely, if you don't want to depend on Bijectors.jl directly, you can just import these functions from their respective packages.
 
@@ -121,7 +121,7 @@ end
 ```
 
 !!! warning
-
+    
     This will return `[Inf, Inf]` if `x[3] == 1` (the 'north pole' of the sphere), which may potentially make downstream computations fail. One potential way around this is to add `eps(T)` to the denominator to avoid it ever being zero: you will sometimes see this trick used in Bijectors.jl. However, be aware that the reverse transform has to also be modified accordingly so that the two transforms remain inverses of each other!
 
 When it comes to computing the Jacobian, we find ourselves in a spot of bother.
@@ -171,7 +171,7 @@ $$J = \begin{bmatrix}
 \end{bmatrix}.$$
 
 !!! note
-
+    
     When you see $x_3$ here, don't think 'the variable $x_3$': it's just shorthand for $\pm \sqrt{1 - x_1^2 - x_2^2}$. (And recall that these formulae hold for both choices of sign.)
 
 Its determinant very nicely simplifies to
diff --git a/docs/src/vector.md b/docs/src/vector.md
index 83f90d908..655d7dd4b 100644
--- a/docs/src/vector.md
+++ b/docs/src/vector.md
@@ -7,9 +7,8 @@ It assumes that there are three forms of samples from a distribution `d` that we
  1. **The original form**, which is what `rand(d)` returns.
 
  2. **A vectorised form**, which is a vector that contains a flattened version of the original form.
-
  3. **A linked vectorised form**, which is a vector in which:
-
+    
       + each element is independent; and
       + each element is unconstrained (can take any value in ℝ).
 
diff --git a/src/vector/product/product.jl b/src/vector/product/product.jl
index 1539355fb..dc35b39e9 100644
--- a/src/vector/product/product.jl
+++ b/src/vector/product/product.jl
@@ -54,7 +54,8 @@ end
 for struct_type in (:ProductVecTransform, :ProductVecInvTransform)
     @eval begin
         Base.:(==)(t1::$struct_type, t2::$struct_type) =
-            (t1.transforms == t2.transforms) & (t1.ranges == t2.ranges) &
+            (t1.transforms == t2.transforms) &
+            (t1.ranges == t2.ranges) &
             (t1.base_size == t2.base_size)
         Base.isequal(t1::$struct_type, t2::$struct_type) =
             isequal(t1.transforms, t2.transforms) &&
diff --git a/src/vector/reshaped/reshaped.jl b/src/vector/reshaped/reshaped.jl
index 8cb8e1840..a829406e2 100644
--- a/src/vector/reshaped/reshaped.jl
+++ b/src/vector/reshaped/reshaped.jl
@@ -28,7 +28,8 @@ struct ReshapeWrapper{N1,N2,T1<:NTuple{N1,Int},T2<:NTuple{N2,Int},B}
     bijector::B
 end
 function Base.:(==)(r1::ReshapeWrapper, r2::ReshapeWrapper)
-    return (r1.reshaped_size == r2.reshaped_size) & (r1.original_size == r2.original_size) &
+    return (r1.reshaped_size == r2.reshaped_size) &
+           (r1.original_size == r2.original_size) &
            (r1.bijector == r2.bijector)
 end
 function Base.isequal(r1::ReshapeWrapper, r2::ReshapeWrapper)
@@ -61,7 +62,8 @@ struct InvReshapeWrapper{N1,N2,T1<:NTuple{N1,Int},T2<:NTuple{N2,Int},B}
     inv_bijector::B
 end
 function Base.:(==)(r1::InvReshapeWrapper, r2::InvReshapeWrapper)
-    return (r1.reshaped_size == r2.reshaped_size) & (r1.original_size == r2.original_size) &
+    return (r1.reshaped_size == r2.reshaped_size) &
+           (r1.original_size == r2.original_size) &
            (r1.inv_bijector == r2.inv_bijector)
 end
 function Base.isequal(r1::InvReshapeWrapper, r2::InvReshapeWrapper)
diff --git a/test/ad/enzyme.jl b/test/ad/enzyme.jl
index 29a6e8983..39ef5b177 100644
--- a/test/ad/enzyme.jl
+++ b/test/ad/enzyme.jl
@@ -32,7 +32,8 @@ using Test
         @testset "forward" begin
             # No batches
             @testset for RT in (Const, Duplicated, DuplicatedNoNeed),
-                Tx in (Const, Duplicated), Ty in (Const, Duplicated),
+                Tx in (Const, Duplicated),
+                Ty in (Const, Duplicated),
                 Tz in (Const, Duplicated)
 
                 test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
@@ -40,7 +41,8 @@ using Test
 
             # Batches
             @testset for RT in (Const, BatchDuplicated, BatchDuplicatedNoNeed),
-                Tx in (Const, BatchDuplicated), Ty in (Const, BatchDuplicated),
+                Tx in (Const, BatchDuplicated),
+                Ty in (Const, BatchDuplicated),
                 Tz in (Const, BatchDuplicated)
 
                 test_forward(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))
@@ -49,7 +51,8 @@ using Test
         @testset "reverse" begin
             # No batches
             @testset for RT in (Const, Active),
-                Tx in (Const, Active), Ty in (Const, Active),
+                Tx in (Const, Active),
+                Ty in (Const, Active),
                 Tz in (Const, Active)
 
                 test_reverse(Bijectors.find_alpha, RT, (x, Tx), (y, Ty), (z, Tz))

From 5e3d21e59453451c1563b99143bb03f410319857 Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Sun, 12 Apr 2026 23:58:55 +0100
Subject: [PATCH 06/11] Fix native AD backend edge cases

---
 ext/BijectorsEnzymeExt.jl      | 110 ++++++++++++++++++++++++++-------
 ext/BijectorsForwardDiffExt.jl |  24 +++++--
 ext/BijectorsMooncakeExt.jl    |  88 +++++++++++++++++++-------
 ext/BijectorsReverseDiffExt.jl |   6 ++
 4 files changed, 177 insertions(+), 51 deletions(-)

diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl
index ccf880f46..f7a7a2eec 100644
--- a/ext/BijectorsEnzymeExt.jl
+++ b/ext/BijectorsEnzymeExt.jl
@@ -5,22 +5,47 @@ import ADTypes: AutoEnzyme
 using Enzyme: Enzyme
 using EnzymeCore: EnzymeCore
 
-# When mode is unspecified, fall back to Reverse.
-function _enzyme_mode(::AutoEnzyme{Nothing})
-    return Enzyme.set_runtime_activity(Enzyme.Reverse)
+const AnyFunctionDuplicated = Union{
+    EnzymeCore.Duplicated,EnzymeCore.DuplicatedNoNeed,EnzymeCore.MixedDuplicated
+}
+
+# `AutoEnzyme()` leaves mode selection to the operation, matching the old DI-backed
+# behaviour: reverse mode for scalar gradients, forward mode for Jacobians.
+_gradient_mode(backend::AutoEnzyme{<:EnzymeCore.ForwardMode}) = backend.mode
+_gradient_mode(backend::AutoEnzyme{<:EnzymeCore.ReverseMode}) = backend.mode
+_gradient_mode(::AutoEnzyme{Nothing}) = Enzyme.Reverse
+
+_forward_withprimal_mode(backend::AutoEnzyme) = Enzyme.WithPrimal(_gradient_mode(backend))
+
+function _gradient_withprimal_mode(backend::AutoEnzyme{<:EnzymeCore.ReverseMode})
+    return EnzymeCore.WithPrimal(backend.mode)
 end
-function _enzyme_mode(backend::AutoEnzyme)
-    return Enzyme.set_runtime_activity(backend.mode)
+_gradient_withprimal_mode(::AutoEnzyme{Nothing}) = Enzyme.ReverseWithPrimal
+
+_jacobian_mode(backend::AutoEnzyme) = backend.mode
+_jacobian_mode(::AutoEnzyme{Nothing}) = Enzyme.Forward
+_jacobian_withprimal_mode(backend::AutoEnzyme) = Enzyme.WithPrimal(_jacobian_mode(backend))
+
+_enzyme_annotate_f(f, ::AutoEnzyme{M,Nothing}, mode) where {M} = f
+_enzyme_annotate_f(f, ::AutoEnzyme{M,<:EnzymeCore.Const}, mode) where {M} = Enzyme.Const(f)
+function _enzyme_duplicated_function(f, ::Type{<:EnzymeCore.MixedDuplicated})
+    return Enzyme.Duplicated(f, Enzyme.make_zero(f))
+end
+function _enzyme_duplicated_function(f, ::Type{A}) where {A<:AnyFunctionDuplicated}
+    return A(f, Enzyme.make_zero(f))
 end
 
-# Returns f annotated as requested, or bare f when annotation is Nothing.
-function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}) where {M,A}
-    if A <: EnzymeCore.Const
+function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}, mode) where {M,A<:AnyFunctionDuplicated}
+    # Mutable functors need a duplicated function annotation plus a shadow.
+    if Enzyme.guess_activity(typeof(f), mode) <: EnzymeCore.Const
         return Enzyme.Const(f)
     else
-        return f
+        return _enzyme_duplicated_function(f, A)
     end
 end
+function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}, mode) where {M,A<:EnzymeCore.Annotation}
+    throw(ArgumentError("unsupported Enzyme function annotation $A"))
+end
 
 # For reverse mode (explicit or auto/nothing), use ReverseWithPrimal so value and
 # gradient are computed in a single autodiff call rather than evaluating f twice.
@@ -29,31 +54,68 @@ function _value_and_gradient(
     backend::Union{AutoEnzyme{Nothing},AutoEnzyme{<:EnzymeCore.ReverseMode}},
     x::AbstractVector,
 )
-    annotated_f = _enzyme_annotate_f(f, backend)
+    mode = _gradient_withprimal_mode(backend)
+    annotated_f = _enzyme_annotate_f(f, backend, mode)
     dx = zero(x)
-    _, val = Enzyme.autodiff(
-        Enzyme.set_runtime_activity(Enzyme.ReverseWithPrimal),
-        annotated_f,
-        Enzyme.Active,
-        Enzyme.Duplicated(x, dx),
-    )
+    _, val = Enzyme.autodiff(mode, annotated_f, Enzyme.Active, Enzyme.Duplicated(x, dx))
     return val, dx
 end
 
 # For forward mode the gradient already requires O(n) JVPs; one extra f(x) evaluation
 # is negligible.
 function _value_and_gradient(f, backend::AutoEnzyme, x::AbstractVector)
-    mode = _enzyme_mode(backend)
-    annotated_f = _enzyme_annotate_f(f, backend)
-    grads = Enzyme.gradient(mode, annotated_f, x)
-    return f(x), first(grads)
+    mode = _forward_withprimal_mode(backend)
+    grad = similar(x)
+    fill!(grad, zero(eltype(x)))
+    value = nothing
+    for i in eachindex(x)
+        dx = zero(x)
+        dx[i] = one(eltype(x))
+        directional, primal = Enzyme.autodiff(
+            mode, _enzyme_annotate_f(f, backend, mode), Enzyme.Duplicated(x, dx)
+        )
+        grad[i] = directional
+        if i == firstindex(x)
+            value = primal
+        end
+    end
+    if isnothing(value)
+        value = f(x)
+    end
+    return value, grad
 end
 
 function _value_and_jacobian(f, backend::AutoEnzyme, x::AbstractVector)
-    mode = _enzyme_mode(backend)
-    annotated_f = _enzyme_annotate_f(f, backend)
-    jacs = Enzyme.jacobian(mode, annotated_f, x)
-    return f(x), first(jacs)
+    mode = _jacobian_mode(backend)
+    if mode isa EnzymeCore.ReverseMode
+        annotated_f = _enzyme_annotate_f(f, backend, mode)
+        jacs = Enzyme.jacobian(mode, annotated_f, x)
+        value = f(x)
+        return value, reshape(first(jacs), length(value), length(x))
+    end
+
+    withprimal_mode = _jacobian_withprimal_mode(backend)
+    value = nothing
+    J = nothing
+    for i in eachindex(x)
+        dx = zero(x)
+        dx[i] = one(eltype(x))
+        directional, primal = Enzyme.autodiff(
+            withprimal_mode,
+            _enzyme_annotate_f(f, backend, withprimal_mode),
+            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(value)
+        value = f(x)
+        J = Matrix{eltype(value)}(undef, length(value), 0)
+    end
+    return value, J
 end
 
 end
diff --git a/ext/BijectorsForwardDiffExt.jl b/ext/BijectorsForwardDiffExt.jl
index 60075b7d9..c71726d81 100644
--- a/ext/BijectorsForwardDiffExt.jl
+++ b/ext/BijectorsForwardDiffExt.jl
@@ -4,16 +4,32 @@ import Bijectors: Bijectors, find_alpha, _value_and_gradient, _value_and_jacobia
 import ADTypes: AutoForwardDiff
 using ForwardDiff: ForwardDiff
 
-function _value_and_gradient(f, ::AutoForwardDiff, x::AbstractVector)
+function _value_and_gradient(
+    f, backend::AutoForwardDiff{chunksize,T}, x::AbstractVector
+) where {chunksize,T}
     result = ForwardDiff.DiffResults.GradientResult(x)
-    ForwardDiff.gradient!(result, f, 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, ::AutoForwardDiff, x::AbstractVector)
+function _value_and_jacobian(
+    f, backend::AutoForwardDiff{chunksize,T}, x::AbstractVector
+) where {chunksize,T}
     y = f(x)
     result = ForwardDiff.DiffResults.JacobianResult(y, x)
-    ForwardDiff.jacobian!(result, f, 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
 
diff --git a/ext/BijectorsMooncakeExt.jl b/ext/BijectorsMooncakeExt.jl
index 0225e2b64..f2873657f 100644
--- a/ext/BijectorsMooncakeExt.jl
+++ b/ext/BijectorsMooncakeExt.jl
@@ -9,31 +9,57 @@ using Mooncake:
     tangent_type,
     @from_chainrules,
     prepare_pullback_cache,
+    prepare_gradient_cache,
     prepare_derivative_cache,
     value_and_pullback!!,
+    value_and_gradient!!,
     value_and_derivative!!,
-    zero_tangent
+    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, ::AutoMooncake, x::AbstractVector{T}) where {T}
-    cache = prepare_pullback_cache(f, x)
-    val, (_, x_grad) = value_and_pullback!!(cache, one(T), f, x)
+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, ::AutoMooncake, x::AbstractVector{T}) where {T}
-    y = f(x)
-    n_out, n_in = length(y), length(x)
-    J = Matrix{T}(undef, n_out, n_in)
-    cache = prepare_pullback_cache(f, x)
-    dy = zeros(eltype(y), n_out)
-    for i in 1:n_out
-        fill!(dy, zero(eltype(y)))
-        dy[i] = one(eltype(y))
+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
@@ -42,14 +68,22 @@ end
 
 ## Forward-mode implementations (column-by-column JVPs)
 
-function _value_and_gradient(f, ::AutoMooncakeForward, x::AbstractVector{T}) where {T}
+function _value_and_gradient(
+    f, backend::AutoMooncakeForward, x::AbstractVector{T}
+) where {T}
     val = f(x)
     n = length(x)
-    grad = Vector{T}(undef, n)
-    cache = prepare_derivative_cache(f, x)
-    df = zero_tangent(f)
+    cache = prepare_derivative_cache(f, x; config=_mooncake_config(backend))
+    df = _mooncake_zero_tangent_or_primal(f, backend)
+    if n == 0
+        return val, Vector{eltype(x)}(undef, 0)
+    end
     dx = zeros(T, n)
-    for j in 1:n
+    dx[1] = one(T)
+    _, first_jvp = value_and_derivative!!(cache, (f, df), (x, dx))
+    grad = Vector{typeof(first_jvp)}(undef, n)
+    grad[1] = first_jvp
+    for j in 2:n
         fill!(dx, zero(T))
         dx[j] = one(T)
         _, jvp = value_and_derivative!!(cache, (f, df), (x, dx))
@@ -58,14 +92,22 @@ function _value_and_gradient(f, ::AutoMooncakeForward, x::AbstractVector{T}) whe
     return val, grad
 end
 
-function _value_and_jacobian(f, ::AutoMooncakeForward, x::AbstractVector{T}) where {T}
+function _value_and_jacobian(
+    f, backend::AutoMooncakeForward, x::AbstractVector{T}
+) where {T}
     y = f(x)
     n_out, n_in = length(y), length(x)
-    J = Matrix{T}(undef, n_out, n_in)
-    cache = prepare_derivative_cache(f, x)
-    df = zero_tangent(f)
+    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)
-    for j in 1: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))
diff --git a/ext/BijectorsReverseDiffExt.jl b/ext/BijectorsReverseDiffExt.jl
index 4f56e3dd4..565cbd922 100644
--- a/ext/BijectorsReverseDiffExt.jl
+++ b/ext/BijectorsReverseDiffExt.jl
@@ -14,6 +14,9 @@ import Bijectors: _value_and_gradient, _value_and_jacobian
 import ADTypes: AutoReverseDiff
 
 function _value_and_gradient(f, ::AutoReverseDiff{false}, x::AbstractVector)
+    if isempty(x)
+        return f(x), similar(x, 0)
+    end
     result = ReverseDiff.DiffResults.GradientResult(x)
     ReverseDiff.gradient!(result, f, x)
     return ReverseDiff.DiffResults.value(result), ReverseDiff.DiffResults.gradient(result)
@@ -27,6 +30,9 @@ function _value_and_jacobian(f, ::AutoReverseDiff{false}, x::AbstractVector)
 end
 
 function _value_and_gradient(f, ::AutoReverseDiff{true}, x::AbstractVector)
+    if isempty(x)
+        return f(x), similar(x, 0)
+    end
     tape = ReverseDiff.GradientTape(f, x)
     compiled = ReverseDiff.compile(tape)
     result = ReverseDiff.DiffResults.GradientResult(x)

From 9ccf137eec7d9d4133d72cd11e1de7229acbb712 Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Mon, 13 Apr 2026 00:05:09 +0100
Subject: [PATCH 07/11] Use Mooncake value_and_gradient!! helpers

---
 Project.toml                |  2 +-
 ext/BijectorsMooncakeExt.jl | 26 ++++----------------------
 test/Project.toml           |  2 +-
 3 files changed, 6 insertions(+), 24 deletions(-)

diff --git a/Project.toml b/Project.toml
index 6cee9a60c..39dfdfc80 100644
--- a/Project.toml
+++ b/Project.toml
@@ -63,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"
diff --git a/ext/BijectorsMooncakeExt.jl b/ext/BijectorsMooncakeExt.jl
index f2873657f..5e7739e5a 100644
--- a/ext/BijectorsMooncakeExt.jl
+++ b/ext/BijectorsMooncakeExt.jl
@@ -68,28 +68,10 @@ end
 
 ## Forward-mode implementations (column-by-column JVPs)
 
-function _value_and_gradient(
-    f, backend::AutoMooncakeForward, x::AbstractVector{T}
-) where {T}
-    val = f(x)
-    n = length(x)
-    cache = prepare_derivative_cache(f, x; config=_mooncake_config(backend))
-    df = _mooncake_zero_tangent_or_primal(f, backend)
-    if n == 0
-        return val, Vector{eltype(x)}(undef, 0)
-    end
-    dx = zeros(T, n)
-    dx[1] = one(T)
-    _, first_jvp = value_and_derivative!!(cache, (f, df), (x, dx))
-    grad = Vector{typeof(first_jvp)}(undef, n)
-    grad[1] = first_jvp
-    for j in 2:n
-        fill!(dx, zero(T))
-        dx[j] = one(T)
-        _, jvp = value_and_derivative!!(cache, (f, df), (x, dx))
-        grad[j] = jvp
-    end
-    return val, grad
+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(
diff --git a/test/Project.toml b/test/Project.toml
index aeb288b9b..5cecc35d8 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -53,7 +53,7 @@ LazyArrays = "1, 2"
 LogDensityProblems = "2"
 LogExpFunctions = "0.3.1"
 MCMCDiagnosticTools = "0.3"
-Mooncake = "0.4, 0.5"
+Mooncake = "0.5.26"
 PDMats = "0.11"
 ReverseDiff = "1.4.2"
 StableRNGs = "1"

From 669d4a82c21b2ba2763a19f48d6dd25c243eb2e2 Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Mon, 13 Apr 2026 00:11:40 +0100
Subject: [PATCH 08/11] Remove DifferentiationInterface references

---
 docs/Project.toml                    |  1 -
 docs/src/defining_examples.md        |  4 +---
 test/Project.toml                    |  4 ++--
 test/integration/enzyme/Project.toml |  1 -
 test/integration/enzyme/main.jl      |  7 +++----
 test/runtests.jl                     |  8 ++++----
 test/vector/cholesky.jl              | 18 +++++++++---------
 test/vector/matrix.jl                |  4 ++--
 test/vector/order.jl                 | 10 +++++-----
 test/vector/product.jl               | 18 +++++++++---------
 test/vector/reshaped.jl              | 12 ++++++------
 11 files changed, 41 insertions(+), 46 deletions(-)

diff --git a/docs/Project.toml b/docs/Project.toml
index 9698442af..6e0d207a8 100644
--- a/docs/Project.toml
+++ b/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"
diff --git a/docs/src/defining_examples.md b/docs/src/defining_examples.md
index 95556b844..4c706853b 100644
--- a/docs/src/defining_examples.md
+++ b/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]
 ```
 
diff --git a/test/Project.toml b/test/Project.toml
index 5cecc35d8..5c7a0a872 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,4 +1,5 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
@@ -7,7 +8,6 @@ ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
@@ -32,6 +32,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [compat]
+ADTypes = "1"
 AbstractMCMC = "5"
 AbstractPPL = "0.14"
 AdvancedHMC = "0.6, 0.7, 0.8"
@@ -39,7 +40,6 @@ ChainRules = "1"
 ChainRulesTestUtils = "0.7, 1"
 ChangesOfVariables = "0.1"
 Combinatorics = "1.0.2"
-DifferentiationInterface = "0.7.7"
 Distributions = "0.25"
 DistributionsAD = "0.6.3"
 Documenter = "1"
diff --git a/test/integration/enzyme/Project.toml b/test/integration/enzyme/Project.toml
index 3ce74b7ab..fe9714f77 100644
--- a/test/integration/enzyme/Project.toml
+++ b/test/integration/enzyme/Project.toml
@@ -1,7 +1,6 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
diff --git a/test/integration/enzyme/main.jl b/test/integration/enzyme/main.jl
index 45cb5c5be..45a414a1a 100644
--- a/test/integration/enzyme/main.jl
+++ b/test/integration/enzyme/main.jl
@@ -1,17 +1,16 @@
 using ADTypes
 using Bijectors
-using DifferentiationInterface
 using Enzyme
 using FiniteDifferences
 using LinearAlgebra
 using Test
 
-const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1))
+const REF_FDM = central_fdm(5, 1)
 
 function test_ad(f, backend, x; rtol=1e-6, atol=1e-6)
     @info "testing AD for function $f with $backend"
-    ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x)
-    gradient = DifferentiationInterface.gradient(f, backend, x)
+    ref_gradient = only(FiniteDifferences.grad(REF_FDM, f, x))
+    _, gradient = Bijectors._value_and_gradient(f, backend, x)
     @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol)
 end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 97074c797..06fa45dcc 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,8 +1,8 @@
 using Bijectors
 
+using ADTypes
 using ChainRulesTestUtils
 using Combinatorics
-using DifferentiationInterface
 using DistributionsAD
 using Documenter: Documenter
 using FiniteDifferences
@@ -57,12 +57,12 @@ if TEST_ENZYME
         ),
     ]
 end
-const REF_BACKEND = AutoFiniteDifferences(; fdm=central_fdm(5, 1))
+const REF_FDM = central_fdm(5, 1)
 
 function test_ad(f, backend, x; rtol=1e-6, atol=1e-6)
     @info "testing AD for function $f with $backend"
-    ref_gradient = DifferentiationInterface.gradient(f, REF_BACKEND, x)
-    gradient = DifferentiationInterface.gradient(f, backend, x)
+    ref_gradient = only(FiniteDifferences.grad(REF_FDM, f, x))
+    _, gradient = Bijectors._value_and_gradient(f, backend, x)
     @test isapprox(gradient, ref_gradient; rtol=rtol, atol=atol)
 end
 
diff --git a/test/vector/cholesky.jl b/test/vector/cholesky.jl
index 74ea1fd41..c11eb93ed 100644
--- a/test/vector/cholesky.jl
+++ b/test/vector/cholesky.jl
@@ -3,8 +3,8 @@ module VBCholeskyTests
 using Distributions
 using LinearAlgebra
 using Test
+using ADTypes: AutoEnzyme, AutoMooncake, AutoMooncakeForward, AutoReverseDiff
 using Bijectors.VectorBijectors
-import DifferentiationInterface as DI
 using ForwardDiff: ForwardDiff
 using ReverseDiff: ReverseDiff
 using Mooncake: Mooncake
@@ -13,17 +13,17 @@ using Enzyme: Enzyme, set_runtime_activity, Const, Forward, Reverse
 # Need runtime activity for some reason.
 # TODO(penelopeysm): Report upstream
 const adtypes = [
-    DI.AutoReverseDiff(),
-    DI.AutoReverseDiff(; compile=true),
-    DI.AutoMooncake(),
-    DI.AutoMooncakeForward(),
-    DI.AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const),
-    DI.AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const),
+    AutoReverseDiff(),
+    AutoReverseDiff(; compile=true),
+    AutoMooncake(),
+    AutoMooncakeForward(),
+    AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const),
+    AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const),
 ]
 
 dists = [
-    # Note: can't test LKJCholesky(1, ...) because its linked vector is length-zero and
-    # DifferentiationInterface trips up with empty vectors.
+    LKJCholesky(1, 1.0, 'U'),
+    LKJCholesky(1, 1.0, 'L'),
     LKJCholesky(3, 1.0, 'U'),
     LKJCholesky(3, 1.0, 'L'),
     LKJCholesky(5, 1.0, 'U'),
diff --git a/test/vector/matrix.jl b/test/vector/matrix.jl
index 53c5c8109..30d1f3c6e 100644
--- a/test/vector/matrix.jl
+++ b/test/vector/matrix.jl
@@ -4,7 +4,7 @@ using Distributions
 using LinearAlgebra
 using Test
 using PDMats
-import DifferentiationInterface as DI
+using ADTypes: AutoMooncake, AutoMooncakeForward
 using Bijectors.VectorBijectors
 using Enzyme: Enzyme
 using ForwardDiff: ForwardDiff
@@ -19,7 +19,7 @@ M = [1 2 3; 4 5 6]
 # TODO(penelopeysm): ReverseDiff gives wrong results when differentiating
 # through VecCorrBijector. Correctness tests are disabled for now.
 # https://github.com/TuringLang/Bijectors.jl/issues/434
-lkj_test_adtypes = [DI.AutoMooncake(), DI.AutoMooncakeForward()]
+lkj_test_adtypes = [AutoMooncake(), AutoMooncakeForward()]
 
 # Don't check that from_linked_vec(d)(randn(...)) is in support for LKJ,
 # The reason is because the inverse bijector for LKJ causes the diagonal
diff --git a/test/vector/order.jl b/test/vector/order.jl
index 682bb1f0d..9ae058e34 100644
--- a/test/vector/order.jl
+++ b/test/vector/order.jl
@@ -3,9 +3,9 @@ module VBOrderTests
 using Distributions
 using LinearAlgebra
 using Test
+using ADTypes: AutoEnzyme, AutoMooncake, AutoMooncakeForward
 using Bijectors.VectorBijectors
 using Bijectors: ordered
-import DifferentiationInterface as DI
 using Enzyme: Enzyme
 using ForwardDiff: ForwardDiff
 using ReverseDiff: ReverseDiff
@@ -24,10 +24,10 @@ base_dists = [
 # because of the heavy setindex! usage.
 # https://github.com/JuliaDiff/ReverseDiff.jl/issues/43 We just avoid testing it for now.
 joint_test_adtypes = [
-    DI.AutoMooncake(),
-    DI.AutoMooncakeForward(),
-    DI.AutoEnzyme(; mode=Enzyme.Forward),
-    DI.AutoEnzyme(; mode=Enzyme.Reverse),
+    AutoMooncake(),
+    AutoMooncakeForward(),
+    AutoEnzyme(; mode=Enzyme.Forward),
+    AutoEnzyme(; mode=Enzyme.Reverse),
 ]
 
 @testset "Order statistics" begin
diff --git a/test/vector/product.jl b/test/vector/product.jl
index ecf5f680e..2137bcc76 100644
--- a/test/vector/product.jl
+++ b/test/vector/product.jl
@@ -4,28 +4,28 @@ using Distributions
 using LinearAlgebra
 using FillArrays: Fill
 using Test
+using ADTypes: AutoEnzyme, AutoMooncake, AutoMooncakeForward, AutoReverseDiff
 using Bijectors.VectorBijectors
-import DifferentiationInterface as DI
 using ForwardDiff: ForwardDiff
 using ReverseDiff: ReverseDiff
 using Mooncake: Mooncake
 using Enzyme: Enzyme, set_runtime_activity, Const, Forward, Reverse
 
 adtypes = [
-    DI.AutoReverseDiff(),
-    DI.AutoReverseDiff(; compile=true),
-    DI.AutoMooncake(),
-    DI.AutoMooncakeForward(),
+    AutoReverseDiff(),
+    AutoReverseDiff(; compile=true),
+    AutoMooncake(),
+    AutoMooncakeForward(),
     # Need runtime activity for some reason.
     # TODO(penelopeysm): Report upstream
-    DI.AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const),
-    DI.AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const),
+    AutoEnzyme(; mode=set_runtime_activity(Forward), function_annotation=Const),
+    AutoEnzyme(; mode=set_runtime_activity(Reverse), function_annotation=Const),
 ]
 
 # Enzyme segfaults on 1.12 + Windows.
 # https://github.com/EnzymeAD/Enzyme.jl/issues/2986
 if VERSION >= v"1.12-" && Sys.iswindows()
-    filter!(a -> !(a isa DI.AutoEnzyme), adtypes)
+    filter!(a -> !(a isa AutoEnzyme), adtypes)
 end
 
 # These are purposely chosen because the vec_length output is the same but
@@ -116,7 +116,7 @@ enzyme_failures = [
         )
     end
 
-    no_enzyme_adtypes = filter(adtype -> !(adtype isa DI.AutoEnzyme), adtypes)
+    no_enzyme_adtypes = filter(adtype -> !(adtype isa AutoEnzyme), adtypes)
     for d in enzyme_failures
         VectorBijectors.test_all(
             d;
diff --git a/test/vector/reshaped.jl b/test/vector/reshaped.jl
index 64d07ce99..9674d2fce 100644
--- a/test/vector/reshaped.jl
+++ b/test/vector/reshaped.jl
@@ -3,8 +3,8 @@ module VBReshapedTests
 using Distributions
 using LinearAlgebra
 using Test
+using ADTypes: AutoEnzyme, AutoMooncake, AutoMooncakeForward, AutoReverseDiff
 using Bijectors.VectorBijectors
-import DifferentiationInterface as DI
 using Enzyme: Enzyme
 using ForwardDiff: ForwardDiff
 using ReverseDiff: ReverseDiff
@@ -30,11 +30,11 @@ reshaped = [
 
 # Fails on 1.10: https://github.com/EnzymeAD/Enzyme.jl/issues/2987
 adtypes_no_enz_rvs = [
-    DI.AutoReverseDiff(),
-    DI.AutoReverseDiff(; compile=true),
-    DI.AutoMooncake(),
-    DI.AutoMooncakeForward(),
-    DI.AutoEnzyme(; mode=Enzyme.Forward, function_annotation=Enzyme.Const),
+    AutoReverseDiff(),
+    AutoReverseDiff(; compile=true),
+    AutoMooncake(),
+    AutoMooncakeForward(),
+    AutoEnzyme(; mode=Enzyme.Forward, function_annotation=Enzyme.Const),
 ]
 reshaped_no_enzyme = [reshape(Beta(2, 2), (1, 1, 1, 1, 1))]
 

From a468093851c2f89fbf5449adfdd6c00eb507d928 Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Mon, 13 Apr 2026 00:37:27 +0100
Subject: [PATCH 09/11] Clean extension imports and AD edge cases

---
 ext/BijectorsEnzymeExt.jl      | 118 ++++++++++++++-------------------
 ext/BijectorsForwardDiffExt.jl |   6 ++
 ext/BijectorsLazyArraysExt.jl  |   8 +--
 ext/BijectorsMooncakeExt.jl    |   7 +-
 ext/BijectorsReverseDiffExt.jl |   6 +-
 5 files changed, 66 insertions(+), 79 deletions(-)

diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl
index f7a7a2eec..09bed6ce4 100644
--- a/ext/BijectorsEnzymeExt.jl
+++ b/ext/BijectorsEnzymeExt.jl
@@ -5,114 +5,98 @@ import ADTypes: AutoEnzyme
 using Enzyme: Enzyme
 using EnzymeCore: EnzymeCore
 
-const AnyFunctionDuplicated = Union{
+const DuplicatedFunctionAnnotations = Union{
     EnzymeCore.Duplicated,EnzymeCore.DuplicatedNoNeed,EnzymeCore.MixedDuplicated
 }
 
-# `AutoEnzyme()` leaves mode selection to the operation, matching the old DI-backed
-# behaviour: reverse mode for scalar gradients, forward mode for Jacobians.
-_gradient_mode(backend::AutoEnzyme{<:EnzymeCore.ForwardMode}) = backend.mode
-_gradient_mode(backend::AutoEnzyme{<:EnzymeCore.ReverseMode}) = backend.mode
-_gradient_mode(::AutoEnzyme{Nothing}) = Enzyme.Reverse
-
-_forward_withprimal_mode(backend::AutoEnzyme) = Enzyme.WithPrimal(_gradient_mode(backend))
-
-function _gradient_withprimal_mode(backend::AutoEnzyme{<:EnzymeCore.ReverseMode})
-    return EnzymeCore.WithPrimal(backend.mode)
-end
-_gradient_withprimal_mode(::AutoEnzyme{Nothing}) = Enzyme.ReverseWithPrimal
-
-_jacobian_mode(backend::AutoEnzyme) = backend.mode
-_jacobian_mode(::AutoEnzyme{Nothing}) = Enzyme.Forward
-_jacobian_withprimal_mode(backend::AutoEnzyme) = Enzyme.WithPrimal(_jacobian_mode(backend))
-
-_enzyme_annotate_f(f, ::AutoEnzyme{M,Nothing}, mode) where {M} = f
-_enzyme_annotate_f(f, ::AutoEnzyme{M,<:EnzymeCore.Const}, mode) where {M} = Enzyme.Const(f)
-function _enzyme_duplicated_function(f, ::Type{<:EnzymeCore.MixedDuplicated})
-    return Enzyme.Duplicated(f, Enzyme.make_zero(f))
-end
-function _enzyme_duplicated_function(f, ::Type{A}) where {A<:AnyFunctionDuplicated}
-    return A(f, Enzyme.make_zero(f))
-end
-
-function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}, mode) where {M,A<:AnyFunctionDuplicated}
-    # Mutable functors need a duplicated function annotation plus a shadow.
-    if Enzyme.guess_activity(typeof(f), mode) <: EnzymeCore.Const
+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
-        return _enzyme_duplicated_function(f, A)
+        throw(ArgumentError("unsupported Enzyme function annotation $annotation"))
     end
 end
-function _enzyme_annotate_f(f, ::AutoEnzyme{M,A}, mode) where {M,A<:EnzymeCore.Annotation}
-    throw(ArgumentError("unsupported Enzyme function annotation $A"))
-end
 
-# For reverse mode (explicit or auto/nothing), use ReverseWithPrimal so value and
-# gradient are computed in a single autodiff call rather than evaluating f twice.
 function _value_and_gradient(
     f,
     backend::Union{AutoEnzyme{Nothing},AutoEnzyme{<:EnzymeCore.ReverseMode}},
     x::AbstractVector,
 )
-    mode = _gradient_withprimal_mode(backend)
-    annotated_f = _enzyme_annotate_f(f, backend, mode)
+    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
 
-# For forward mode the gradient already requires O(n) JVPs; one extra f(x) evaluation
-# is negligible.
-function _value_and_gradient(f, backend::AutoEnzyme, x::AbstractVector)
-    mode = _forward_withprimal_mode(backend)
-    grad = similar(x)
-    fill!(grad, zero(eltype(x)))
-    value = nothing
+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, _enzyme_annotate_f(f, backend, mode), Enzyme.Duplicated(x, dx)
-        )
+        directional, primal = Enzyme.autodiff(mode, annotated_f, Enzyme.Duplicated(x, dx))
         grad[i] = directional
         if i == firstindex(x)
             value = primal
         end
     end
-    if isnothing(value)
-        value = f(x)
-    end
     return value, grad
 end
 
-function _value_and_jacobian(f, backend::AutoEnzyme, x::AbstractVector)
-    mode = _jacobian_mode(backend)
-    if mode isa EnzymeCore.ReverseMode
-        annotated_f = _enzyme_annotate_f(f, backend, mode)
-        jacs = Enzyme.jacobian(mode, annotated_f, x)
-        value = f(x)
-        return value, reshape(first(jacs), length(value), length(x))
+function _value_and_jacobian(
+    f, backend::AutoEnzyme{<:EnzymeCore.ReverseMode}, x::AbstractVector
+)
+    value = f(x)
+    if isempty(x)
+        return value, Matrix{eltype(value)}(undef, length(value), 0)
     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
 
-    withprimal_mode = _jacobian_withprimal_mode(backend)
-    value = nothing
+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(
-            withprimal_mode,
-            _enzyme_annotate_f(f, backend, withprimal_mode),
-            Enzyme.Duplicated(x, dx),
-        )
+        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(value)
-        value = f(x)
+    if isnothing(J)
         J = Matrix{eltype(value)}(undef, length(value), 0)
     end
     return value, J
diff --git a/ext/BijectorsForwardDiffExt.jl b/ext/BijectorsForwardDiffExt.jl
index c71726d81..4a4f9f27c 100644
--- a/ext/BijectorsForwardDiffExt.jl
+++ b/ext/BijectorsForwardDiffExt.jl
@@ -7,6 +7,9 @@ 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
@@ -22,6 +25,9 @@ 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
diff --git a/ext/BijectorsLazyArraysExt.jl b/ext/BijectorsLazyArraysExt.jl
index 03943fdc7..6be9502a5 100644
--- a/ext/BijectorsLazyArraysExt.jl
+++ b/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
 
diff --git a/ext/BijectorsMooncakeExt.jl b/ext/BijectorsMooncakeExt.jl
index 5e7739e5a..6f2af0216 100644
--- a/ext/BijectorsMooncakeExt.jl
+++ b/ext/BijectorsMooncakeExt.jl
@@ -1,11 +1,10 @@
 module BijectorsMooncakeExt
 
+using Mooncake: Mooncake
 using Mooncake:
     @is_primitive,
     MinimalCtx,
-    Mooncake,
     CoDual,
-    primal,
     tangent_type,
     @from_chainrules,
     prepare_pullback_cache,
@@ -144,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()
diff --git a/ext/BijectorsReverseDiffExt.jl b/ext/BijectorsReverseDiffExt.jl
index 565cbd922..1a54d2c96 100644
--- a/ext/BijectorsReverseDiffExt.jl
+++ b/ext/BijectorsReverseDiffExt.jl
@@ -50,14 +50,12 @@ function _value_and_jacobian(f, ::AutoReverseDiff{true}, x::AbstractVector)
 end
 
 using Bijectors:
-    ChainRulesCore,
     Elementwise,
     SimplexBijector,
     maphcat,
     simplex_link_jacobian,
     simplex_invlink_jacobian,
-    simplex_logabsdetjac_gradient,
-    Inverse
+    simplex_logabsdetjac_gradient
 import Bijectors:
     Bijectors,
     _eps,
@@ -66,8 +64,6 @@ import Bijectors:
     _simplex_bijector,
     _simplex_inv_bijector,
     replace_diag,
-    jacobian,
-    _inv_link_chol_lkj,
     _link_chol_lkj,
     _transform_ordered,
     _transform_inverse_ordered,

From 81bbaa421c92e0867ad4508000893f84e7b672ec Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Mon, 13 Apr 2026 01:00:31 +0100
Subject: [PATCH 10/11] Handle empty Enzyme reverse Jacobians

---
 ext/BijectorsEnzymeExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl
index 09bed6ce4..e89b2147a 100644
--- a/ext/BijectorsEnzymeExt.jl
+++ b/ext/BijectorsEnzymeExt.jl
@@ -67,7 +67,7 @@ function _value_and_jacobian(
     f, backend::AutoEnzyme{<:EnzymeCore.ReverseMode}, x::AbstractVector
 )
     value = f(x)
-    if isempty(x)
+    if isempty(x) || isempty(value)
         return value, Matrix{eltype(value)}(undef, length(value), 0)
     end
     annotated_f = _annotate_function(f, backend, backend.mode)

From e4d8585a2479de60708af2a0dea8ff929fb1df81 Mon Sep 17 00:00:00 2001
From: Hong Ge <hg344@cam.ac.uk>
Date: Mon, 13 Apr 2026 01:20:54 +0100
Subject: [PATCH 11/11] Fix empty Enzyme Jacobian output shape

---
 ext/BijectorsEnzymeExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/BijectorsEnzymeExt.jl b/ext/BijectorsEnzymeExt.jl
index e89b2147a..af1ce5638 100644
--- a/ext/BijectorsEnzymeExt.jl
+++ b/ext/BijectorsEnzymeExt.jl
@@ -68,7 +68,7 @@ function _value_and_jacobian(
 )
     value = f(x)
     if isempty(x) || isempty(value)
-        return value, Matrix{eltype(value)}(undef, length(value), 0)
+        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))