Use AbstractPPL AD interface in tests.

CLAUDE.md

@@ -15,7 +15,7 @@ The code in `src/` outside of `src/vector/` is largely legacy. It still works an
 
 ## Testing
 
-Run tests with `julia --project -e 'using Pkg; Pkg.test()'`.
+Run tests with `julia --project -e 'using Pkg; Pkg.test()'`. The Enzyme integration suite lives in `test/integration_tests/enzyme/` and runs separately (see `.github/workflows/Enzyme.yml`). Always refresh each environment (`Pkg.update()` / `up`) before tests or doc builds — a stale manifest can cause subtle resolution and loading issues.
 
 `Bijectors.VectorBijectors.test_all(d)` is the single source of truth for whether a VectorBijectors implementation is correct. If it passes, the implementation is almost certainly fine. It checks roundtrips, type stability, vector lengths, optics, allocations, log-Jacobian correctness (against ForwardDiff), and AD compatibility across backends.
 

Project.toml

@@ -4,9 +4,9 @@ version = "0.15.24"
 
 [deps]
 AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
-DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
@@ -42,12 +42,12 @@ BijectorsReverseDiffChainRulesExt = ["ChainRules", "ReverseDiff"]
 BijectorsReverseDiffExt = "ReverseDiff"
 
 [compat]
-AbstractPPL = "0.14"
+AbstractPPL = "0.15"
+ADTypes = "1"
 ArgCheck = "1, 2"
 ChainRules = "1"
 ChainRulesCore = "0.10.11, 1"
 ChangesOfVariables = "0.1"
-DifferentiationInterface = "0.7.14"
 Distributions = "0.25.33"
 DocStringExtensions = "0.9"
 EnzymeCore = "0.8.15"

docs/Project.toml

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

docs/src/defining_examples.md

@@ -199,6 +199,8 @@ If we try to calculate a Jacobian for `StereographicProj()`, we will just get a
 So, we need to take an extra step to map from the independent coordinates of $x$ (i.e., $x_1$ and $x_2$) to the full 3D coordinates, and _then_ to the plane:
 
 ```@example stereographic
+using ForwardDiff, LinearAlgebra
+
 sgn = 1
 
 function full_transform(x12)
@@ -207,12 +209,8 @@ 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 = ForwardDiff.jacobian(full_transform, x[1:2])
 logjac = logabsdet(jac)[1]
 ```
 

docs/src/vector.md

@@ -117,7 +117,7 @@ For more information about generally testing bijectors (and in particular how to
 
 One of the most tricky parts of testing Bijectors is ensuring that the transforms are compatible with automatic differentiation.
 This is important for DynamicPPL: we need to be able to compute the gradient of the log-density with respect to (possibly transformed) parameters, which may include the log-abs-det-Jacobian of the transformation.
-`test_all` accepts an `adtypes` keyword argument and defaults to `[AutoForwardDiff()]`, so the main test suite exercises ForwardDiff as the reference backend.
+`test_all` accepts an `adtypes` keyword argument that defaults to an empty list, so the main test suite does not exercise any AD backend directly — ForwardDiff is still used internally as the reference for Jacobian correctness via `ref_adtype`.
 ReverseDiff, Mooncake, and Enzyme are covered by standalone integration suites under `test/integration_tests/<backend>/`, each of which passes its own `adtypes` list through to `test_all`.
 It is acceptable to skip tests for a particular backend if there are genuine upstream bugs, especially with ReverseDiff, which is not actively maintained.
 Where possible it is best to ensure that all backends are supported, and to use `@test_broken` to mark any known issues with specific backends.

src/vector/test_utils.jl

@@ -1,14 +1,14 @@
 using Test
 using LinearAlgebra: logabsdet, Cholesky, UpperTriangular, LowerTriangular
-import DifferentiationInterface as DI
-
-# 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()
+using ADTypes
+using AbstractPPL: AbstractPPL
 
 # ForwardDiff is the baked-in reference; per-backend integration suites under
-# `test/integration_tests/` pass their own `adtypes` to `test_all`.
-const default_adtypes = [ref_adtype]
+# `test/integration_tests/` pass their own `adtypes` to `test_all`. Routing the reference
+# through `AbstractPPL.prepare(ref_adtype, ...)` keeps the call path uniform with the
+# backend AD calls.
+const ref_adtype = AutoForwardDiff()
+const default_adtypes = ADTypes.AbstractADType[]
 
 _get_value_support(::D.Distribution{<:Any,VS}) where {VS<:D.ValueSupport} = VS
 
@@ -228,7 +228,7 @@ function test_all(
     d::D.Distribution;
     expected_zero_allocs=(),
     adtypes=default_adtypes,
-    broken_adtypes=DI.AbstractADType[],
+    broken_adtypes=ADTypes.AbstractADType[],
     ad_atol=1e-10,
     ad_rtol=sqrt(eps()),
     roundtrip_atol=1e-10,
@@ -409,6 +409,10 @@ function test_optics(d::D.Distribution)
         end
     end
 
+    # Discrete: link transform is identity, so the AD Jacobian sparsity test below is
+    # trivial — skip it.
+    d isa D.Distribution{<:Any,D.Discrete} && return nothing
+
     @testset "linked_optic_vec: $(_name(d))" begin
         # This is a lot harder to test. What we need to prove is that
         #       x = rand(d)
@@ -422,7 +426,10 @@ 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)
+        f = to_linked_vec(d) ∘ from_vec(d)
+        J = last(
+            AbstractPPL.value_and_jacobian!!(AbstractPPL.prepare(ref_adtype, f, xvec), xvec)
+        )
         o = optic_vec(d)
         lo = linked_optic_vec(d)
         for i in 1:length(yvec)
@@ -537,6 +544,10 @@ function test_logjac(d::D.Distribution, atol, rtol)
         end
     end
 
+    # Discrete: link transform is identity, so the linked logjac is trivially zero on
+    # both sides of the comparison — skip the AD-based check.
+    d isa D.Distribution{<:Any,D.Discrete} && return nothing
+
     # Link logjacs will not be zero, so we need to check against a chosen backend. Because
     # Jacobians need to map from vector to vector, here we test the transformation of the
     # vectorised form to the linked vectorised form via the original sample.
@@ -566,7 +577,15 @@ 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(
+                            AbstractPPL.value_and_jacobian!!(
+                                AbstractPPL.prepare(ref_adtype, ad_ffwd, ad_xvec), ad_xvec
+                            ),
+                        ),
+                    ),
+                )
                 @test vbt_logjac ≈ ad_logjac atol = atol rtol = rtol
             end
 
@@ -579,7 +598,15 @@ 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(
+                            AbstractPPL.value_and_jacobian!!(
+                                AbstractPPL.prepare(ref_adtype, ad_frvs, yvec), yvec
+                            ),
+                        ),
+                    ),
+                )
                 @test vbt_logjac ≈ ad_logjac atol = atol rtol = rtol
             end
         end
@@ -590,26 +617,42 @@ 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{<:ADTypes.AbstractADType}, atol, rtol)
     # Discrete-distribution link transforms are identities; nothing AD-meaningful to check.
     d isa D.Distribution{<:Any,D.Discrete} && return nothing
 
     @testset "AD forward: $(_name(d))" begin
         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 = last(
+            AbstractPPL.value_and_jacobian!!(
+                AbstractPPL.prepare(ref_adtype, ffwd, xvec), xvec
+            ),
+        )
 
         ladj(xvec) = last(with_logabsdet_jacobian(ffwd, xvec))
-        ref_grad_ladj = DI.gradient(ladj, ref_adtype, xvec)
+        ref_grad_ladj = last(
+            AbstractPPL.value_and_gradient!!(
+                AbstractPPL.prepare(ref_adtype, ladj, xvec), xvec
+            ),
+        )
 
         for adtype in adtypes
             @testset let x = x, adtype = adtype, d = d
-                ad_jac = DI.jacobian(ffwd, adtype, xvec)
+                ad_jac = last(
+                    AbstractPPL.value_and_jacobian!!(
+                        AbstractPPL.prepare(adtype, ffwd, xvec), 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 = last(
+                    AbstractPPL.value_and_gradient!!(
+                        AbstractPPL.prepare(adtype, ladj, xvec), xvec
+                    ),
+                )
                 @test ref_grad_ladj ≈ ad_grad_ladj atol = atol rtol = rtol
             end
         end
@@ -619,19 +662,35 @@ 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 = last(
+            AbstractPPL.value_and_jacobian!!(
+                AbstractPPL.prepare(ref_adtype, frvs, yvec), yvec
+            ),
+        )
 
         ladj(yvec) = last(with_logabsdet_jacobian(frvs, yvec))
-        ref_grad_ladj = DI.gradient(ladj, ref_adtype, yvec)
+        ref_grad_ladj = last(
+            AbstractPPL.value_and_gradient!!(
+                AbstractPPL.prepare(ref_adtype, ladj, yvec), yvec
+            ),
+        )
 
         for adtype in adtypes
             @testset let x = x, adtype = adtype, d = d
-                ad_jac = DI.jacobian(frvs, adtype, yvec)
+                ad_jac = last(
+                    AbstractPPL.value_and_jacobian!!(
+                        AbstractPPL.prepare(adtype, frvs, yvec), 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 = last(
+                    AbstractPPL.value_and_gradient!!(
+                        AbstractPPL.prepare(adtype, ladj, yvec), yvec
+                    ),
+                )
                 @test ref_grad_ladj ≈ ad_grad_ladj atol = atol rtol = rtol
             end
         end
@@ -641,7 +700,7 @@ end
 # Broken-mode counterpart to `test_ad`: collapses the four per-adtype AD comparisons into
 # one `@test_broken` per adtype so a partially-broken backend stays cleanly broken.
 function _test_ad_broken(
-    d::D.Distribution, adtypes::Vector{<:DI.AbstractADType}, atol, rtol
+    d::D.Distribution, adtypes::Vector{<:ADTypes.AbstractADType}, atol, rtol
 )
     d isa D.Distribution{<:Any,D.Discrete} && return nothing
     @testset "AD (broken): $(_name(d))" begin
@@ -652,27 +711,65 @@ function _test_ad_broken(
         frvs = to_vec(d) ∘ from_linked_vec(d)
         ladj_fwd(xv) = last(with_logabsdet_jacobian(ffwd, xv))
         ladj_rev(yv) = last(with_logabsdet_jacobian(frvs, yv))
-        ref_jac_fwd = DI.jacobian(ffwd, ref_adtype, xvec)
-        ref_grad_ladj_fwd = DI.gradient(ladj_fwd, ref_adtype, xvec)
-        ref_jac_rev = DI.jacobian(frvs, ref_adtype, yvec)
-        ref_grad_ladj_rev = DI.gradient(ladj_rev, ref_adtype, yvec)
+        ref_jac_fwd = last(
+            AbstractPPL.value_and_jacobian!!(
+                AbstractPPL.prepare(ref_adtype, ffwd, xvec), xvec
+            ),
+        )
+        ref_grad_ladj_fwd = last(
+            AbstractPPL.value_and_gradient!!(
+                AbstractPPL.prepare(ref_adtype, ladj_fwd, xvec), xvec
+            ),
+        )
+        ref_jac_rev = last(
+            AbstractPPL.value_and_jacobian!!(
+                AbstractPPL.prepare(ref_adtype, frvs, yvec), yvec
+            ),
+        )
+        ref_grad_ladj_rev = last(
+            AbstractPPL.value_and_gradient!!(
+                AbstractPPL.prepare(ref_adtype, ladj_rev, yvec), yvec
+            ),
+        )
         for adtype in adtypes
             @testset let x = x, adtype = adtype, d = d
                 @test_broken (
                     isapprox(
-                        DI.jacobian(ffwd, adtype, xvec), ref_jac_fwd; atol=atol, rtol=rtol
+                        last(
+                            AbstractPPL.value_and_jacobian!!(
+                                AbstractPPL.prepare(adtype, ffwd, xvec), xvec
+                            ),
+                        ),
+                        ref_jac_fwd;
+                        atol=atol,
+                        rtol=rtol,
                     ) &&
                     isapprox(
-                        DI.gradient(ladj_fwd, adtype, xvec),
+                        last(
+                            AbstractPPL.value_and_gradient!!(
+                                AbstractPPL.prepare(adtype, ladj_fwd, xvec), xvec
+                            ),
+                        ),
                         ref_grad_ladj_fwd;
                         atol=atol,
                         rtol=rtol,
                     ) &&
                     isapprox(
-                        DI.jacobian(frvs, adtype, yvec), ref_jac_rev; atol=atol, rtol=rtol
+                        last(
+                            AbstractPPL.value_and_jacobian!!(
+                                AbstractPPL.prepare(adtype, frvs, yvec), yvec
+                            ),
+                        ),
+                        ref_jac_rev;
+                        atol=atol,
+                        rtol=rtol,
                     ) &&
                     isapprox(
-                        DI.gradient(ladj_rev, adtype, yvec),
+                        last(
+                            AbstractPPL.value_and_gradient!!(
+                                AbstractPPL.prepare(adtype, ladj_rev, yvec), yvec
+                            ),
+                        ),
                         ref_grad_ladj_rev;
                         atol=atol,
                         rtol=rtol,

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"
@@ -27,15 +28,16 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 
 [compat]
+ADTypes = "1"
 AbstractMCMC = "5"
-AbstractPPL = "0.14"
+AbstractPPL = "0.15"
 AdvancedHMC = "0.6, 0.7, 0.8"
 Bijectors = "0.15"
 ChainRules = "1"
 ChainRulesTestUtils = "0.7, 1"
 ChangesOfVariables = "0.1"
 Combinatorics = "1.0.2"
-DifferentiationInterface = "0.7.7"
+DifferentiationInterface = "0.6, 0.7"
 Distributions = "0.25"
 FillArrays = "1"
 FiniteDifferences = "0.11, 0.12"

test/integration_tests/enzyme/Project.toml

@@ -1,12 +1,12 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+AbstractPPL = "7a57a42e-76ec-4ea3-a279-07e840d6d9cf"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
-FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
@@ -17,4 +17,6 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Bijectors = {path = "../../.."}
 
 [compat]
+AbstractPPL = "0.15"
+DifferentiationInterface = "0.6, 0.7"
 Enzyme = "0.13.111"