Repository: TuringLang/DynamicPPL.jl
Pull request #882: Implement AD testing and benchmarking (hand rolled)
Author: login: penelopeysm
Created: 2025-04-04T00:57:40Z | Merged: 2025-04-14T15:05:36Z

One of two options. The other one at #883.

This PR implements functionality for testing _and_ benchmarking AD. It is largely copied over from [my ModelTests repo](https://github.com/penelopeysm/ModelTests.jl) where I've been playing around with this.

Closes #869

## What does it contain?

It basically adds one function `DynamicPPL.TestUtils.AD.run_ad`. See the docstring for more info.

## Why not an extension?

The only new dependencies are Statistics, which is stdlib, and Chairmarks, which itself has no non-stdlib dependencies. I therefore consider it unnecessary to add an extension (which would bring a number of drawbacks, e.g. reduced discoverability as users have to load the trigger packages themselves, us having to faff around with functions declared in `src/` and extended in `ext/`, ...)

## Why do I like this one more?

See #883.

Linked issue #869: AD testing
I'm not aware of an issue for this, so wanted to open one to capture thoughts.

The code for this is still not in DynamicPPL, it's hosted at https://github.com/penelopeysm/ModelTests.jl as it's easier for me to iterate on it there

Desiderata:

1. Each AD backend runs in its own CI.
2. For each AD backend, each model tested runs in its own process. This is pretty awkward. I think it basically means we need to have a shell script calling Julia.
3. The result of the job should be aggregated - if any model fails then the job should have a red cross
4. Output should specify benchmark time (if run successfully) & error (if not). When the jobs finish running, this info must be collated into a single csv and/or html on gh-pages i.e. this info must be easily available to end-user

Note: some of these are _difficult_ to do right. It may well be that we should sacrifice some of these points, or push them to later, just for the sake of getting _something_ out.

Bonus stretch goals:

1. Avoid recalculating the 'ground truth' with ForwardDiff for the same model multiple times.
2. Add links to existing GitHub issues when reasons for failing models are known.
3. Add ability to test on different varinfos

Additional details in https://github.com/TuringLang/DynamicPPL.jl/pull/799#issuecomment-2643772496

Diff:
diff --git a/HISTORY.md b/HISTORY.md
index f43f543c3..52616315d 100644
--- a/HISTORY.md
+++ b/HISTORY.md
@@ -1,5 +1,10 @@
 # DynamicPPL Changelog
 
+## 0.35.8
+
+Added the `DynamicPPL.TestUtils.AD.run_ad` function to test the correctness and/or benchmark the performance of an automatic differentiation backend on DynamicPPL models.
+Please see [the docstring](https://turinglang.org/DynamicPPL.jl/api/#DynamicPPL.TestUtils.AD.run_ad) for more information.
+
 ## 0.35.7
 
 `check_model_and_trace` now errors if any NaN's are encountered when evaluating the model.
diff --git a/Project.toml b/Project.toml
index 5790eaeb1..6c72e55d9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.35.7"
+version = "0.35.8"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -10,6 +10,7 @@ Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 BangBang = "198e06fe-97b7-11e9-32a5-e1d131e6ad66"
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
@@ -22,6 +23,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [weakdeps]
@@ -49,6 +51,7 @@ Accessors = "0.1"
 BangBang = "0.4.1"
 Bijectors = "0.13.18, 0.14, 0.15"
 ChainRulesCore = "1"
+Chairmarks = "1.3.1"
 Compat = "4"
 ConstructionBase = "1.5.4"
 DifferentiationInterface = "0.6.41"
@@ -67,5 +70,6 @@ Mooncake = "0.4.95"
 OrderedCollections = "1"
 Random = "1.6"
 Requires = "1"
+Statistics = "1"
 Test = "1.6"
 julia = "1.10"
diff --git a/docs/src/api.md b/docs/src/api.md
index 9c8249c97..7af4074e2 100644
--- a/docs/src/api.md
+++ b/docs/src/api.md
@@ -205,7 +205,16 @@ values_as_in_model
 NamedDist
 ```
 
-## Testing Utilities
+## AD testing and benchmarking utilities
+
+To test and/or benchmark the performance of an AD backend on a model, DynamicPPL provides the following utilities:
+
+```@docs
+DynamicPPL.TestUtils.AD.run_ad
+DynamicPPL.TestUtils.AD.ADResult
+```
+
+## Demo models
 
 DynamicPPL provides several demo models and helpers for testing samplers in the `DynamicPPL.TestUtils` submodule.
 
diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl
index 50fe0edc7..6e26d9a97 100644
--- a/src/DynamicPPL.jl
+++ b/src/DynamicPPL.jl
@@ -175,7 +175,6 @@ include("context_implementations.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")
 include("submodel_macro.jl")
-include("test_utils.jl")
 include("transforming.jl")
 include("logdensityfunction.jl")
 include("model_utils.jl")
@@ -184,6 +183,7 @@ include("values_as_in_model.jl")
 
 include("debug_utils.jl")
 using .DebugUtils
+include("test_utils.jl")
 
 include("experimental.jl")
 include("deprecated.jl")
diff --git a/src/test_utils.jl b/src/test_utils.jl
index c7d12c927..65079f023 100644
--- a/src/test_utils.jl
+++ b/src/test_utils.jl
@@ -18,5 +18,6 @@ include("test_utils/models.jl")
 include("test_utils/contexts.jl")
 include("test_utils/varinfo.jl")
 include("test_utils/sampler.jl")
+include("test_utils/ad.jl")
 
 end
diff --git a/src/test_utils/ad.jl b/src/test_utils/ad.jl
new file mode 100644
index 000000000..06c76df5e
--- /dev/null
+++ b/src/test_utils/ad.jl
@@ -0,0 +1,209 @@
+module AD
+
+using ADTypes: AbstractADType, AutoForwardDiff
+using Chairmarks: @be
+import DifferentiationInterface as DI
+using DocStringExtensions
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using LogDensityProblems: logdensity, logdensity_and_gradient
+using Random: Random, Xoshiro
+using Statistics: median
+using Test: @test
+
+export ADResult, run_ad
+
+# This function needed to work around the fact that different backends can
+# return different AbstractArrays for the gradient. See
+# https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754 for more
+# context.
+_to_vec_f64(x::AbstractArray) = x isa Vector{Float64} ? x : collect(Float64, x)
+
+"""
+    REFERENCE_ADTYPE
+
+Reference AD backend to use for comparison. In this case, ForwardDiff.jl, since
+it's the default AD backend used in Turing.jl.
+"""
+const REFERENCE_ADTYPE = AutoForwardDiff()
+
+"""
+    ADResult
+
+Data structure to store the results of the AD correctness test.
+"""
+struct ADResult
+    "The DynamicPPL model that was tested"
+    model::Model
+    "The VarInfo that was used"
+    varinfo::AbstractVarInfo
+    "The values at which the model was evaluated"
+    params::Vector{<:Real}
+    "The AD backend that was tested"
+    adtype::AbstractADType
+    "The absolute tolerance for the value of logp"
+    value_atol::Real
+    "The absolute tolerance for the gradient of logp"
+    grad_atol::Real
+    "The expected value of logp"
+    value_expected::Union{Nothing,Float64}
+    "The expected gradient of logp"
+    grad_expected::Union{Nothing,Vector{Float64}}
+    "The value of logp (calculated using `adtype`)"
+    value_actual::Union{Nothing,Real}
+    "The gradient of logp (calculated using `adtype`)"
+    grad_actual::Union{Nothing,Vector{Float64}}
+    "If benchmarking was requested, the time taken by the AD backend to calculate the gradient of logp, divided by the time taken to evaluate logp itself"
+    time_vs_primal::Union{Nothing,Float64}
+end
+
+"""
+    run_ad(
+        model::Model,
+        adtype::ADTypes.AbstractADType;
+        test=true,
+        benchmark=false,
+        value_atol=1e-6,
+        grad_atol=1e-6,
+        varinfo::AbstractVarInfo=VarInfo(model),
+        params::Vector{<:Real}=varinfo[:],
+        reference_adtype::ADTypes.AbstractADType=REFERENCE_ADTYPE,
+        expected_value_and_grad::Union{Nothing,Tuple{Real,Vector{<:Real}}}=nothing,
+        verbose=true,
+    )::ADResult
+
+Test the correctness and/or benchmark the AD backend `adtype` for the model
+`model`.
+
+Whether to test and benchmark is controlled by the `test` and `benchmark`
+keyword arguments. By default, `test` is `true` and `benchmark` is `false`.
+
+Returns an [`ADResult`](@ref) object, which contains the results of the
+test and/or benchmark.
+
+Note that to run AD successfully you will need to import the AD backend itself.
+For example, to test with `AutoReverseDiff()` you will need to run `import
+ReverseDiff`.
+
+There are two positional arguments, which absolutely must be provided:
+
+1. `model` - The model being tested.
+2. `adtype` - The AD backend being tested.
+
+Everything else is optional, and can be categorised into several groups:
+
+1. _How to specify the VarInfo._
+
+   DynamicPPL contains several different types of VarInfo objects which change
+   the way model evaluation occurs. If you want to use a specific type of
+   VarInfo, pass it as the `varinfo` argument. Otherwise, it will default to
+   using a `TypedVarInfo` generated from the model.
+
+2. _How to specify the parameters._
+
+   For maximum control over this, generate a vector of parameters yourself and
+   pass this as the `params` argument. If you don't specify this, it will be
+   taken from the contents of the VarInfo.
+
+   Note that if the VarInfo is not specified (and thus automatically generated)
+   the parameters in it will have been sampled from the prior of the model. If
+   you want to seed the parameter generation, the easiest way is to pass a
+   `rng` argument to the VarInfo constructor (i.e. do `VarInfo(rng, model)`).
+
+   Finally, note that these only reflect the parameters used for _evaluating_
+   the gradient. If you also want to control the parameters used for
+   _preparing_ the gradient, then you need to manually set these parameters in
+   the VarInfo object, for example using `vi = DynamicPPL.unflatten(vi,
+   prep_params)`. You could then evaluate the gradient at a different set of
+   parameters using the `params` keyword argument.
+
+3. _How to specify the results to compare against._ (Only if `test=true`.)
+
+   Once logp and its gradient has been calculated with the specified `adtype`,
+   it must be tested for correctness.
+
+   This can be done either by specifying `reference_adtype`, in which case logp
+   and its gradient will also be calculated with this reference in order to
+   obtain the ground truth; or by using `expected_value_and_grad`, which is a
+   tuple of `(logp, gradient)` that the calculated values must match. The
+   latter is useful if you are testing multiple AD backends and want to avoid
+   recalculating the ground truth multiple times.
+
+   The default reference backend is ForwardDiff. If none of these parameters are
+   specified, ForwardDiff will be used to calculate the ground truth.
+
+4. _How to specify the tolerances._ (Only if `test=true`.)
+
+   The tolerances for the value and gradient can be set using `value_atol` and
+   `grad_atol`. These default to 1e-6.
+
+5. _Whether to output extra logging information._
+
+   By default, this function prints messages when it runs. To silence it, set
+   `verbose=false`.
+"""
+function run_ad(
+    model::Model,
+    adtype::AbstractADType;
+    test=true,
+    benchmark=false,
+    value_atol=1e-6,
+    grad_atol=1e-6,
+    varinfo::AbstractVarInfo=VarInfo(model),
+    params::Vector{<:Real}=varinfo[:],
+    reference_adtype::AbstractADType=REFERENCE_ADTYPE,
+    expected_value_and_grad::Union{Nothing,Tuple{Real,Vector{<:Real}}}=nothing,
+    verbose=true,
+)::ADResult
+    verbose && @info "Running AD on $(model.f) with $(adtype)\n"
+    params = map(identity, params)
+    verbose && println("       params : $(params)")
+    ldf = LogDensityFunction(model, varinfo; adtype=adtype)
+
+    value, grad = logdensity_and_gradient(ldf, params)
+    grad = _to_vec_f64(grad)
+    verbose && println("       actual : $((value, grad))")
+
+    if test
+        # Calculate ground truth to compare against
+        value_true, grad_true = if expected_value_and_grad === nothing
+            ldf_reference = LogDensityFunction(model, varinfo; adtype=reference_adtype)
+            logdensity_and_gradient(ldf_reference, params)
+        else
+            expected_value_and_grad
+        end
+        verbose && println("     expected : $((value_true, grad_true))")
+        grad_true = _to_vec_f64(grad_true)
+        # Then compare
+        @test isapprox(value, value_true; atol=value_atol)
+        @test isapprox(grad, grad_true; atol=grad_atol)
+    else
+        value_true = nothing
+        grad_true = nothing
+    end
+
+    time_vs_primal = if benchmark
+        primal_benchmark = @be (ldf, params) logdensity(_[1], _[2])
+        grad_benchmark = @be (ldf, params) logdensity_and_gradient(_[1], _[2])
+        t = median(grad_benchmark).time / median(primal_benchmark).time
+        verbose && println("grad / primal : $(t)")
+        t
+    else
+        nothing
+    end
+
+    return ADResult(
+        model,
+        varinfo,
+        params,
+        adtype,
+        value_atol,
+        grad_atol,
+        value_true,
+        grad_true,
+        value,
+        grad,
+        time_vs_primal,
+    )
+end
+
+end # module DynamicPPL.TestUtils.AD
diff --git a/test/ad.jl b/test/ad.jl
index a4f3dbfa7..33d581228 100644
--- a/test/ad.jl
+++ b/test/ad.jl
@@ -56,10 +56,12 @@ using DynamicPPL: LogDensityFunction
                             ref_ldf, adtype
                         )
                     else
-                        ldf = DynamicPPL.LogDensityFunction(ref_ldf, adtype)
-                        logp, grad = LogDensityProblems.logdensity_and_gradient(ldf, x)
-                        @test grad ≈ ref_grad
-                        @test logp ≈ ref_logp
+                        DynamicPPL.TestUtils.AD.run_ad(
+                            m,
+                            adtype;
+                            varinfo=varinfo,
+                            expected_value_and_grad=(ref_logp, ref_grad),
+                        )
                     end
                 end
             end


Reviews:
user: login: sunxd3
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T12:24:35Z
user: login: penelopeysm
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T14:17:54Z
user: login: penelopeysm
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-08T14:20:21Z
user: login: yebai
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-14T14:02:19Z
user: login: penelopeysm
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-14T14:04:40Z

Review comments:
diff_hunk: @@ -0,0 +1,190 @@
+module AD
+
+using ADTypes: AbstractADType, AutoForwardDiff
+using Chairmarks: @be
+import DifferentiationInterface as DI
+using DocStringExtensions
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using LogDensityProblems: logdensity, logdensity_and_gradient
+using Random: Random, Xoshiro
+using Statistics: median
+using Test: @test
+
+export ADResult, run_ad
+
+# This function needed to work around the fact that different backends can
+# return different AbstractArrays for the gradient. See
+# https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754 for more
+# context.
+_to_vec_f64(x::AbstractArray) = x isa Vector{Float64} ? x : collect(Float64, x)
+
+"""
+    REFERENCE_ADTYPE
+
+Reference AD backend to use for comparison. In this case, ForwardDiff.jl, since
+it's the default AD backend used in Turing.jl.
+"""
+const REFERENCE_ADTYPE = AutoForwardDiff()
+
+"""
+    ADResult
+
+Data structure to store the results of the AD correctness test.
+
+If you want to quickly check whether the result is a success or failure, you
+can use `isnothing(result.error)`.
+"""
+struct ADResult
+    "The DynamicPPL model that was tested"
+    model::Model
+    "The VarInfo that was used"
+    varinfo::AbstractVarInfo
+    "The values at which the model was evaluated"
+    params::Vector{<:Real}
+    "The AD backend that was tested"
+    adtype::AbstractADType
+    "The absolute tolerance for the value of logp"
+    value_atol::Real
+    "The absolute tolerance for the gradient of logp"
+    grad_atol::Real
+    "The expected value of logp"
+    value_expected::Union{Nothing,Float64}
+    "The expected gradient of logp"
+    grad_expected::Union{Nothing,Vector{Float64}}
+    "The value of logp (calculated using `adtype`)"
+    value_actual::Union{Nothing,Real}
+    "The gradient of logp (calculated using `adtype`)"
+    grad_actual::Union{Nothing,Vector{Float64}}
+    "If benchmarking was requested, the time taken by the AD backend to calculate the gradient of logp, divided by the time taken to evaluate logp itself"
+    time_vs_primal::Union{Nothing,Float64}
+end
+
+"""
+    run_ad(
+        model::Model,
+        adtype::ADTypes.AbstractADType;
+        test=true,
+        benchmark=false,
+        value_atol=1e-6,
+        grad_atol=1e-6,
+        varinfo::AbstractVarInfo=VarInfo(model),
+        params::Vector{<:Real}=varinfo[:],
+        reference_adtype::ADTypes.AbstractADType=REFERENCE_ADTYPE,
+        expected_value_and_grad::Union{Nothing,Tuple{Real,Vector{<:Real}}}=nothing,
+        verbose=true,
+    )::ADResult
+
+Test the correctness and/or benchmark the AD backend `adtype` for the model
+`model`.
+
+Whether to test and benchmark is controlled by the `test` and `benchmark`
+keyword arguments. By default, `test` is `true` and `benchmark` is `false.
path: src/test_utils/ad.jl
user: login: sunxd3
body: ```suggestion
keyword arguments. By default, `test` is `true` and `benchmark` is `false`.
```
created_at: 2025-04-08T11:58:25Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,190 @@
+module AD
+
+using ADTypes: AbstractADType, AutoForwardDiff
+using Chairmarks: @be
+import DifferentiationInterface as DI
+using DocStringExtensions
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using LogDensityProblems: logdensity, logdensity_and_gradient
+using Random: Random, Xoshiro
+using Statistics: median
+using Test: @test
+
+export ADResult, run_ad
+
+# This function needed to work around the fact that different backends can
+# return different AbstractArrays for the gradient. See
+# https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754 for more
+# context.
+_to_vec_f64(x::AbstractArray) = x isa Vector{Float64} ? x : collect(Float64, x)
+
+"""
+    REFERENCE_ADTYPE
+
+Reference AD backend to use for comparison. In this case, ForwardDiff.jl, since
+it's the default AD backend used in Turing.jl.
+"""
+const REFERENCE_ADTYPE = AutoForwardDiff()
+
+"""
+    ADResult
+
+Data structure to store the results of the AD correctness test.
+
+If you want to quickly check whether the result is a success or failure, you
+can use `isnothing(result.error)`.
+"""
+struct ADResult
+    "The DynamicPPL model that was tested"
+    model::Model
+    "The VarInfo that was used"
+    varinfo::AbstractVarInfo
+    "The values at which the model was evaluated"
+    params::Vector{<:Real}
+    "The AD backend that was tested"
+    adtype::AbstractADType
+    "The absolute tolerance for the value of logp"
+    value_atol::Real
+    "The absolute tolerance for the gradient of logp"
+    grad_atol::Real
+    "The expected value of logp"
+    value_expected::Union{Nothing,Float64}
+    "The expected gradient of logp"
+    grad_expected::Union{Nothing,Vector{Float64}}
+    "The value of logp (calculated using `adtype`)"
+    value_actual::Union{Nothing,Real}
+    "The gradient of logp (calculated using `adtype`)"
+    grad_actual::Union{Nothing,Vector{Float64}}
+    "If benchmarking was requested, the time taken by the AD backend to calculate the gradient of logp, divided by the time taken to evaluate logp itself"
+    time_vs_primal::Union{Nothing,Float64}
+end
+
+"""
+    run_ad(
+        model::Model,
+        adtype::ADTypes.AbstractADType;
+        test=true,
+        benchmark=false,
+        value_atol=1e-6,
+        grad_atol=1e-6,
+        varinfo::AbstractVarInfo=VarInfo(model),
+        params::Vector{<:Real}=varinfo[:],
+        reference_adtype::ADTypes.AbstractADType=REFERENCE_ADTYPE,
+        expected_value_and_grad::Union{Nothing,Tuple{Real,Vector{<:Real}}}=nothing,
+        verbose=true,
+    )::ADResult
+
+Test the correctness and/or benchmark the AD backend `adtype` for the model
+`model`.
+
+Whether to test and benchmark is controlled by the `test` and `benchmark`
+keyword arguments. By default, `test` is `true` and `benchmark` is `false.
+
+Returns an [`ADResult`](@ref) object, which contains the results of the
+test and/or benchmark.
+
+This function is not as complicated as its signature makes it look. There are
+two things that must be provided:
+
+1. `model` - The model being tested.
+2. `adtype` - The AD backend being tested.
+
+Everything else is optional, and can be categorised into several groups:
+
+1. _How to specify the VarInfo._ DynamicPPL contains several different types of
+VarInfo objects which change the way model evaluation occurs. If you want to
path: src/test_utils/ad.jl
user: login: sunxd3
body: these are messed up a bit (https://turinglang.org/DynamicPPL.jl/previews/PR882/api/#AD-testing-and-benchmarking-utilities) probably need some indentations? 
created_at: 2025-04-08T11:59:39Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,190 @@
+module AD
+
+using ADTypes: AbstractADType, AutoForwardDiff
+using Chairmarks: @be
+import DifferentiationInterface as DI
+using DocStringExtensions
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using LogDensityProblems: logdensity, logdensity_and_gradient
+using Random: Random, Xoshiro
+using Statistics: median
+using Test: @test
+
+export ADResult, run_ad
+
+# This function needed to work around the fact that different backends can
+# return different AbstractArrays for the gradient. See
+# https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754 for more
+# context.
+_to_vec_f64(x::AbstractArray) = x isa Vector{Float64} ? x : collect(Float64, x)
+
+"""
+    REFERENCE_ADTYPE
+
+Reference AD backend to use for comparison. In this case, ForwardDiff.jl, since
+it's the default AD backend used in Turing.jl.
+"""
+const REFERENCE_ADTYPE = AutoForwardDiff()
+
+"""
+    ADResult
+
+Data structure to store the results of the AD correctness test.
+
+If you want to quickly check whether the result is a success or failure, you
+can use `isnothing(result.error)`.
+"""
+struct ADResult
+    "The DynamicPPL model that was tested"
+    model::Model
+    "The VarInfo that was used"
+    varinfo::AbstractVarInfo
+    "The values at which the model was evaluated"
+    params::Vector{<:Real}
+    "The AD backend that was tested"
+    adtype::AbstractADType
+    "The absolute tolerance for the value of logp"
+    value_atol::Real
+    "The absolute tolerance for the gradient of logp"
+    grad_atol::Real
+    "The expected value of logp"
+    value_expected::Union{Nothing,Float64}
+    "The expected gradient of logp"
+    grad_expected::Union{Nothing,Vector{Float64}}
+    "The value of logp (calculated using `adtype`)"
+    value_actual::Union{Nothing,Real}
+    "The gradient of logp (calculated using `adtype`)"
+    grad_actual::Union{Nothing,Vector{Float64}}
+    "If benchmarking was requested, the time taken by the AD backend to calculate the gradient of logp, divided by the time taken to evaluate logp itself"
+    time_vs_primal::Union{Nothing,Float64}
+end
+
+"""
+    run_ad(
+        model::Model,
+        adtype::ADTypes.AbstractADType;
+        test=true,
+        benchmark=false,
+        value_atol=1e-6,
+        grad_atol=1e-6,
+        varinfo::AbstractVarInfo=VarInfo(model),
+        params::Vector{<:Real}=varinfo[:],
+        reference_adtype::ADTypes.AbstractADType=REFERENCE_ADTYPE,
+        expected_value_and_grad::Union{Nothing,Tuple{Real,Vector{<:Real}}}=nothing,
+        verbose=true,
+    )::ADResult
+
+Test the correctness and/or benchmark the AD backend `adtype` for the model
+`model`.
+
+Whether to test and benchmark is controlled by the `test` and `benchmark`
+keyword arguments. By default, `test` is `true` and `benchmark` is `false.
+
+Returns an [`ADResult`](@ref) object, which contains the results of the
+test and/or benchmark.
+
+This function is not as complicated as its signature makes it look. There are
+two things that must be provided:
+
+1. `model` - The model being tested.
+2. `adtype` - The AD backend being tested.
+
+Everything else is optional, and can be categorised into several groups:
+
+1. _How to specify the VarInfo._ DynamicPPL contains several different types of
+VarInfo objects which change the way model evaluation occurs. If you want to
+use a specific type of VarInfo, pass it as the `varinfo` argument. Otherwise,
+it will default to using a `TypedVarInfo` generated from the model.
+
+2. _How to specify the parameters._ For maximum control over this, generate a
+vector of parameters yourself and pass this as the `params` argument. If you
+don't specify this, it will be taken from the contents of the VarInfo. Note
+that if the VarInfo is not specified (and thus automatically generated) the
+parameters in it will have been sampled from the prior of the model. If you
+want to seed the parameter generation, the easiest way is to pass a `rng`
+argument to the VarInfo constructor (i.e. do `VarInfo(rng, model)`).
+
+3. _How to specify the results to compare against._ (Only if `test=true`.) Once
+logp and its gradient has been calculated with the specified `adtype`, it must
+be tested for correctness. This can be done either by specifying
+`reference_adtype`, in which case logp and its gradient will also be calculated
+with this reference in order to obtain the ground truth; or by using
+`expected_value_and_grad`, which is a tuple of (logp, gradient) that the
+calculated values must match. The latter is useful if you are testing multiple
+AD backends and want to avoid recalculating the ground truth multiple times.
+The default reference backend is ForwardDiff. If none of these parameters are
+specified, that will be used to calculate the ground truth.
path: src/test_utils/ad.jl
user: login: sunxd3
body: maybe mention somewhere that the user would need to install and using `ForwardDiff`? (or any reference backend they want to use?)
created_at: 2025-04-08T12:23:53Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,190 @@
+module AD
+
+using ADTypes: AbstractADType, AutoForwardDiff
+using Chairmarks: @be
+import DifferentiationInterface as DI
+using DocStringExtensions
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using LogDensityProblems: logdensity, logdensity_and_gradient
+using Random: Random, Xoshiro
+using Statistics: median
+using Test: @test
+
+export ADResult, run_ad
+
+# This function needed to work around the fact that different backends can
+# return different AbstractArrays for the gradient. See
+# https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754 for more
+# context.
+_to_vec_f64(x::AbstractArray) = x isa Vector{Float64} ? x : collect(Float64, x)
+
+"""
+    REFERENCE_ADTYPE
+
+Reference AD backend to use for comparison. In this case, ForwardDiff.jl, since
+it's the default AD backend used in Turing.jl.
+"""
+const REFERENCE_ADTYPE = AutoForwardDiff()
+
+"""
+    ADResult
+
+Data structure to store the results of the AD correctness test.
+
+If you want to quickly check whether the result is a success or failure, you
+can use `isnothing(result.error)`.
+"""
+struct ADResult
+    "The DynamicPPL model that was tested"
+    model::Model
+    "The VarInfo that was used"
+    varinfo::AbstractVarInfo
+    "The values at which the model was evaluated"
+    params::Vector{<:Real}
+    "The AD backend that was tested"
+    adtype::AbstractADType
+    "The absolute tolerance for the value of logp"
+    value_atol::Real
+    "The absolute tolerance for the gradient of logp"
+    grad_atol::Real
+    "The expected value of logp"
+    value_expected::Union{Nothing,Float64}
+    "The expected gradient of logp"
+    grad_expected::Union{Nothing,Vector{Float64}}
+    "The value of logp (calculated using `adtype`)"
+    value_actual::Union{Nothing,Real}
+    "The gradient of logp (calculated using `adtype`)"
+    grad_actual::Union{Nothing,Vector{Float64}}
+    "If benchmarking was requested, the time taken by the AD backend to calculate the gradient of logp, divided by the time taken to evaluate logp itself"
+    time_vs_primal::Union{Nothing,Float64}
+end
+
+"""
+    run_ad(
+        model::Model,
+        adtype::ADTypes.AbstractADType;
+        test=true,
+        benchmark=false,
+        value_atol=1e-6,
+        grad_atol=1e-6,
+        varinfo::AbstractVarInfo=VarInfo(model),
+        params::Vector{<:Real}=varinfo[:],
+        reference_adtype::ADTypes.AbstractADType=REFERENCE_ADTYPE,
+        expected_value_and_grad::Union{Nothing,Tuple{Real,Vector{<:Real}}}=nothing,
+        verbose=true,
+    )::ADResult
+
+Test the correctness and/or benchmark the AD backend `adtype` for the model
+`model`.
+
+Whether to test and benchmark is controlled by the `test` and `benchmark`
+keyword arguments. By default, `test` is `true` and `benchmark` is `false.
+
+Returns an [`ADResult`](@ref) object, which contains the results of the
+test and/or benchmark.
+
+This function is not as complicated as its signature makes it look. There are
+two things that must be provided:
+
+1. `model` - The model being tested.
+2. `adtype` - The AD backend being tested.
+
+Everything else is optional, and can be categorised into several groups:
+
+1. _How to specify the VarInfo._ DynamicPPL contains several different types of
+VarInfo objects which change the way model evaluation occurs. If you want to
path: src/test_utils/ad.jl
user: login: penelopeysm
body: 😬 Thanks for spotting
created_at: 2025-04-08T14:17:54Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,190 @@
+module AD
+
+using ADTypes: AbstractADType, AutoForwardDiff
+using Chairmarks: @be
+import DifferentiationInterface as DI
+using DocStringExtensions
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using LogDensityProblems: logdensity, logdensity_and_gradient
+using Random: Random, Xoshiro
+using Statistics: median
+using Test: @test
+
+export ADResult, run_ad
+
+# This function needed to work around the fact that different backends can
+# return different AbstractArrays for the gradient. See
+# https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754 for more
+# context.
+_to_vec_f64(x::AbstractArray) = x isa Vector{Float64} ? x : collect(Float64, x)
+
+"""
+    REFERENCE_ADTYPE
+
+Reference AD backend to use for comparison. In this case, ForwardDiff.jl, since
+it's the default AD backend used in Turing.jl.
+"""
+const REFERENCE_ADTYPE = AutoForwardDiff()
+
+"""
+    ADResult
+
+Data structure to store the results of the AD correctness test.
+
+If you want to quickly check whether the result is a success or failure, you
+can use `isnothing(result.error)`.
+"""
+struct ADResult
+    "The DynamicPPL model that was tested"
+    model::Model
+    "The VarInfo that was used"
+    varinfo::AbstractVarInfo
+    "The values at which the model was evaluated"
+    params::Vector{<:Real}
+    "The AD backend that was tested"
+    adtype::AbstractADType
+    "The absolute tolerance for the value of logp"
+    value_atol::Real
+    "The absolute tolerance for the gradient of logp"
+    grad_atol::Real
+    "The expected value of logp"
+    value_expected::Union{Nothing,Float64}
+    "The expected gradient of logp"
+    grad_expected::Union{Nothing,Vector{Float64}}
+    "The value of logp (calculated using `adtype`)"
+    value_actual::Union{Nothing,Real}
+    "The gradient of logp (calculated using `adtype`)"
+    grad_actual::Union{Nothing,Vector{Float64}}
+    "If benchmarking was requested, the time taken by the AD backend to calculate the gradient of logp, divided by the time taken to evaluate logp itself"
+    time_vs_primal::Union{Nothing,Float64}
+end
+
+"""
+    run_ad(
+        model::Model,
+        adtype::ADTypes.AbstractADType;
+        test=true,
+        benchmark=false,
+        value_atol=1e-6,
+        grad_atol=1e-6,
+        varinfo::AbstractVarInfo=VarInfo(model),
+        params::Vector{<:Real}=varinfo[:],
+        reference_adtype::ADTypes.AbstractADType=REFERENCE_ADTYPE,
+        expected_value_and_grad::Union{Nothing,Tuple{Real,Vector{<:Real}}}=nothing,
+        verbose=true,
+    )::ADResult
+
+Test the correctness and/or benchmark the AD backend `adtype` for the model
+`model`.
+
+Whether to test and benchmark is controlled by the `test` and `benchmark`
+keyword arguments. By default, `test` is `true` and `benchmark` is `false.
+
+Returns an [`ADResult`](@ref) object, which contains the results of the
+test and/or benchmark.
+
+This function is not as complicated as its signature makes it look. There are
+two things that must be provided:
+
+1. `model` - The model being tested.
+2. `adtype` - The AD backend being tested.
+
+Everything else is optional, and can be categorised into several groups:
+
+1. _How to specify the VarInfo._ DynamicPPL contains several different types of
+VarInfo objects which change the way model evaluation occurs. If you want to
+use a specific type of VarInfo, pass it as the `varinfo` argument. Otherwise,
+it will default to using a `TypedVarInfo` generated from the model.
+
+2. _How to specify the parameters._ For maximum control over this, generate a
+vector of parameters yourself and pass this as the `params` argument. If you
+don't specify this, it will be taken from the contents of the VarInfo. Note
+that if the VarInfo is not specified (and thus automatically generated) the
+parameters in it will have been sampled from the prior of the model. If you
+want to seed the parameter generation, the easiest way is to pass a `rng`
+argument to the VarInfo constructor (i.e. do `VarInfo(rng, model)`).
+
+3. _How to specify the results to compare against._ (Only if `test=true`.) Once
+logp and its gradient has been calculated with the specified `adtype`, it must
+be tested for correctness. This can be done either by specifying
+`reference_adtype`, in which case logp and its gradient will also be calculated
+with this reference in order to obtain the ground truth; or by using
+`expected_value_and_grad`, which is a tuple of (logp, gradient) that the
+calculated values must match. The latter is useful if you are testing multiple
+AD backends and want to avoid recalculating the ground truth multiple times.
+The default reference backend is ForwardDiff. If none of these parameters are
+specified, that will be used to calculate the ground truth.
path: src/test_utils/ad.jl
user: login: penelopeysm
body: Added! I have to say DifferentiationInterface has a really nice error message for this too so even if people don't read the docstring they'll be prompted haha
created_at: 2025-04-08T14:20:21Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,212 @@
+module AD
+
+using ADTypes: AbstractADType, AutoForwardDiff
+using Chairmarks: @be
+import DifferentiationInterface as DI
+using DocStringExtensions
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using LogDensityProblems: logdensity, logdensity_and_gradient
+using Random: Random, Xoshiro
+using Statistics: median
+using Test: @test
+
+export ADResult, run_ad
+
+# This function needed to work around the fact that different backends can
+# return different AbstractArrays for the gradient. See
+# https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754 for more
+# context.
+_to_vec_f64(x::AbstractArray) = x isa Vector{Float64} ? x : collect(Float64, x)
+
+"""
+    REFERENCE_ADTYPE
+
+Reference AD backend to use for comparison. In this case, ForwardDiff.jl, since
+it's the default AD backend used in Turing.jl.
+"""
+const REFERENCE_ADTYPE = AutoForwardDiff()
+
+"""
+    ADResult
+
+Data structure to store the results of the AD correctness test.
+
+If you want to quickly check whether the result is a success or failure, you
+can use `isnothing(result.error)`.
path: src/test_utils/ad.jl
user: login: yebai
body: The `ADResult` type doesn't seem to have a field `result.error`? 
created_at: 2025-04-14T14:00:56Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,212 @@
+module AD
+
+using ADTypes: AbstractADType, AutoForwardDiff
+using Chairmarks: @be
+import DifferentiationInterface as DI
+using DocStringExtensions
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using LogDensityProblems: logdensity, logdensity_and_gradient
+using Random: Random, Xoshiro
+using Statistics: median
+using Test: @test
+
+export ADResult, run_ad
+
+# This function needed to work around the fact that different backends can
+# return different AbstractArrays for the gradient. See
+# https://github.com/JuliaDiff/DifferentiationInterface.jl/issues/754 for more
+# context.
+_to_vec_f64(x::AbstractArray) = x isa Vector{Float64} ? x : collect(Float64, x)
+
+"""
+    REFERENCE_ADTYPE
+
+Reference AD backend to use for comparison. In this case, ForwardDiff.jl, since
+it's the default AD backend used in Turing.jl.
+"""
+const REFERENCE_ADTYPE = AutoForwardDiff()
+
+"""
+    ADResult
+
+Data structure to store the results of the AD correctness test.
+
+If you want to quickly check whether the result is a success or failure, you
+can use `isnothing(result.error)`.
path: src/test_utils/ad.jl
user: login: penelopeysm
body: Oh, good spot, I removed it.
created_at: 2025-04-14T14:04:40Z
author_association: MEMBER

PR comments:
user: login: github-actions[bot]
created_at: 2025-04-04T01:06:35Z
body: ## Benchmark Report for Commit b107b92ff61e65e92158c018f897afa9d523444d
### Computer Information
```
Julia Version 1.11.4
Commit 8561cc3d68d (2025-03-10 11:36 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 4 × AMD EPYC 7763 64-Core Processor
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, znver3)
Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores)
```
### Benchmark Results
```
|                 Model | Dimension |  AD Backend |      VarInfo Type | Linked | Eval Time / Ref Time | AD Time / Eval Time |
|-----------------------|-----------|-------------|-------------------|--------|----------------------|---------------------|
| Simple assume observe |         1 | forwarddiff |             typed |  false |                  9.3 |                 1.6 |
|           Smorgasbord |       201 | forwarddiff |             typed |  false |                601.0 |                43.1 |
|           Smorgasbord |       201 | forwarddiff | simple_namedtuple |   true |                423.7 |                46.3 |
|           Smorgasbord |       201 | forwarddiff |           untyped |   true |               1180.2 |                29.1 |
|           Smorgasbord |       201 | forwarddiff |       simple_dict |   true |               3827.5 |                20.5 |
|           Smorgasbord |       201 | reversediff |             typed |   true |               1442.6 |                29.1 |
|           Smorgasbord |       201 |    mooncake |             typed |   true |                919.0 |                 5.3 |
|    Loop univariate 1k |      1000 |    mooncake |             typed |   true |               5402.5 |                 4.1 |
|       Multivariate 1k |      1000 |    mooncake |             typed |   true |               1080.6 |                 8.3 |
|   Loop univariate 10k |     10000 |    mooncake |             typed |   true |              59302.8 |                 3.7 |
|      Multivariate 10k |     10000 |    mooncake |             typed |   true |               8930.4 |                 9.6 |
|               Dynamic |        10 |    mooncake |             typed |   true |                133.8 |                11.8 |
|              Submodel |         1 |    mooncake |             typed |   true |                 25.0 |                 8.0 |
|                   LDA |        12 | reversediff |             typed |   true |                382.9 |                 6.1 |
```
author_association: CONTRIBUTOR
user: login: codecov[bot]
created_at: 2025-04-04T01:09:01Z
body: ## [Codecov](https://app.codecov.io/gh/TuringLang/DynamicPPL.jl/pull/882?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TuringLang) Report
Attention: Patch coverage is `68.96552%` with `9 lines` in your changes missing coverage. Please review.
> Project coverage is 84.80%. Comparing base [(`c7bdc3f`)](https://app.codecov.io/gh/TuringLang/DynamicPPL.jl/commit/c7bdc3f76a47b8cf4f5ca1d8187a14426558b1c2?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TuringLang) to head [(`ef5a1ce`)](https://app.codecov.io/gh/TuringLang/DynamicPPL.jl/commit/ef5a1ceee657f57bd73b5cbd3cba94241e5742f8?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TuringLang).
> Report is 1 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/TuringLang/DynamicPPL.jl/pull/882?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TuringLang) | Patch % | Lines |
|---|---|---|
| [src/test\_utils/ad.jl](https://app.codecov.io/gh/TuringLang/DynamicPPL.jl/pull/882?src=pr&el=tree&filepath=src%2Ftest_utils%2Fad.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TuringLang#diff-c3JjL3Rlc3RfdXRpbHMvYWQuamw=) | 68.96% | [9 Missing :warning: ](https://app.codecov.io/gh/TuringLang/DynamicPPL.jl/pull/882?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TuringLang) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #882      +/-   ##
==========================================
- Coverage   84.92%   84.80%   -0.13%     
==========================================
  Files          34       35       +1     
  Lines        3814     3843      +29     
==========================================
+ Hits         3239     3259      +20     
- Misses        575      584       +9     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/TuringLang/DynamicPPL.jl/pull/882?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TuringLang).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=TuringLang).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: coveralls
created_at: 2025-04-04T01:09:49Z
body: ## Pull Request Test Coverage Report for [Build 14448023153](https://coveralls.io/builds/73215598)


### Details

* **0** of **28**   **(0.0%)**  changed or added relevant lines in **1** file are covered.
* **3** unchanged lines in **1** file lost coverage.
* Overall coverage decreased (**-0.1%**) to **84.892%**

---

|  Changes Missing Coverage | Covered Lines | Changed/Added Lines | % |
| :-----|--------------|--------|---: |
| [src/test_utils/ad.jl](https://coveralls.io/builds/73215598/source?filename=src%2Ftest_utils%2Fad.jl#L19) | 0 | 28 | 0.0%
<!-- | **Total:** | **0** | **28** | **0.0%** | -->

|  Files with Coverage Reduction | New Missed Lines | % |
| :-----|--------------|--: |
| [src/varinfo.jl](https://coveralls.io/builds/73215598/source?filename=src%2Fvarinfo.jl#L323) | 3 | 83.83% |
<!-- | **Total:** | **3** |  | -->

|  Totals | [![Coverage Status](https://coveralls.io/builds/73215598/badge)](https://coveralls.io/builds/73215598) |
| :-- | --: |
| Change from base [Build 14392526425](https://coveralls.io/builds/73176401): |  -0.1% |
| Covered Lines: | 3259 |
| Relevant Lines: | 3839 |

---
##### 💛  - [Coveralls](https://coveralls.io)

author_association: NONE
user: login: yebai
created_at: 2025-04-14T14:03:45Z
body: Thanks, @penelopeysm. This looks good!


author_association: MEMBER