Link varinfo by default in AD testing utilities; make test suite run on linked varinfos

HISTORY.md

@@ -4,6 +4,11 @@
 
 **Breaking changes**
 
+### AD testing utilities
+
+`DynamicPPL.TestUtils.AD.run_ad` now links the VarInfo by default.
+To disable this, pass the `linked=false` keyword argument.
+
 ### VarInfo constructors
 
 `VarInfo(vi::VarInfo, values)` has been removed. You can replace this directly with `unflatten(vi, values)` instead.

src/test_utils/ad.jl

@@ -4,7 +4,7 @@
 using Chairmarks: @be
 import DifferentiationInterface as DI
 using DocStringExtensions
-using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo
+using DynamicPPL: Model, LogDensityFunction, VarInfo, AbstractVarInfo, link
 using LogDensityProblems: logdensity, logdensity_and_gradient
 using Random: Random, Xoshiro
 using Statistics: median
@@ -64,8 +64,9 @@
         benchmark=false,
         value_atol=1e-6,
         grad_atol=1e-6,
+        linked::Bool=true,
         varinfo::AbstractVarInfo=VarInfo(model),
-        params::Vector{<:Real}=varinfo[:],
+        params::Union{Nothing,Vector{<:Real}}=nothing,
         reference_adtype::ADTypes.AbstractADType=REFERENCE_ADTYPE,
         expected_value_and_grad::Union{Nothing,Tuple{Real,Vector{<:Real}}}=nothing,
         verbose=true,
@@ -98,6 +99,11 @@
    VarInfo, pass it as the `varinfo` argument. Otherwise, it will default to
    using a `TypedVarInfo` generated from the model.
 
+   It will also perform _linking_, that is, the parameters in the VarInfo will
+   be transformed to unconstrained Euclidean space if they aren't already in
+   that space. Note that the act of linking may change the length of the
+   parameters. To disable linking, set `linked=false`.
+
 2. _How to specify the parameters._
 
    For maximum control over this, generate a vector of parameters yourself and
@@ -148,14 +154,22 @@
     benchmark=false,
     value_atol=1e-6,
     grad_atol=1e-6,
+    linked::Bool=true,
     varinfo::AbstractVarInfo=VarInfo(model),
-    params::Vector{<:Real}=varinfo[:],
+    params::Union{Nothing,Vector{<:Real}}=nothing,
     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"
+    if linked
+        varinfo = link(varinfo, model)
+    end
+    if isnothing(params)
+        params = varinfo[:]
+    end
     params = map(identity, params)
+
+    verbose && @info "Running AD on $(model.f) with $(adtype)\n"
     verbose && println("       params : $(params)")
     ldf = LogDensityFunction(model, varinfo; adtype=adtype)
 

test/ad.jl

@@ -23,6 +23,7 @@ using DynamicPPL: LogDensityFunction
             varinfos = DynamicPPL.TestUtils.setup_varinfos(m, rand_param_values, vns)
 
             @testset "$(short_varinfo_name(varinfo))" for varinfo in varinfos
+                # TODO: This runs unlinked. Should we test linked as well?
                 f = LogDensityFunction(m, varinfo)
                 x = DynamicPPL.getparams(f)
                 # Calculate reference logp + gradient of logp using ForwardDiff
@@ -56,9 +57,11 @@ using DynamicPPL: LogDensityFunction
                             ref_ldf, adtype
                         )
                     else
+                        # TODO: Should we test linked as well?
                         DynamicPPL.TestUtils.AD.run_ad(
                             m,
                             adtype;
+                            linked=false,
                             varinfo=varinfo,
                             expected_value_and_grad=(ref_logp, ref_grad),
                         )