Test linked varinfos

Closes #891

Author: penelopeysm

HISTORY.md

@@ -12,6 +12,10 @@ If the calculated value or gradient is incorrect, it also throws a `DynamicPPL.T
 This exception contains the actual and expected gradient so you can inspect it if needed; see the documentation for more information.
 From a practical perspective, this means that if you need to add this to a test suite, you need to use `@test run_ad(...) isa Any` rather than just `run_ad(...)`.
 
+### SimpleVarInfo linking / invlinking
+
+Linking a linked SimpleVarInfo, or invlinking an unlinked SimpleVarInfo, now displays a warning instead of an error.
+
 ### VarInfo constructors
 
 `VarInfo(vi::VarInfo, values)` has been removed. You can replace this directly with `unflatten(vi, values)` instead.

src/transforming.jl

@@ -19,9 +19,9 @@ function tilde_assume(
     lp = Bijectors.logpdf_with_trans(right, r, !isinverse)
 
     if istrans(vi, vn)
-        @assert isinverse "Trying to link already transformed variables"
+        isinverse || @warn "Trying to link an already transformed variable ($vn)"
     else
-        @assert !isinverse "Trying to invlink non-transformed variables"
+        isinverse && @warn "Trying to invlink a non-transformed variable ($vn)"
     end
 
     # Only transform if `!isinverse` since `vi[vn, right]`

test/ad.jl

@@ -23,22 +23,23 @@ 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)
+                linked_varinfo = DynamicPPL.link(varinfo, m)
+                f = LogDensityFunction(m, linked_varinfo)
                 x = DynamicPPL.getparams(f)
                 # Calculate reference logp + gradient of logp using ForwardDiff
-                ref_ldf = LogDensityFunction(m, varinfo; adtype=ref_adtype)
+                ref_ldf = LogDensityFunction(m, linked_varinfo; adtype=ref_adtype)
                 ref_logp, ref_grad = LogDensityProblems.logdensity_and_gradient(ref_ldf, x)
 
                 @testset "$adtype" for adtype in test_adtypes
-                    @info "Testing AD on: $(m.f) - $(short_varinfo_name(varinfo)) - $adtype"
+                    @info "Testing AD on: $(m.f) - $(short_varinfo_name(linked_varinfo)) - $adtype"
 
                     # Put predicates here to avoid long lines
                     is_mooncake = adtype isa AutoMooncake
                     is_1_10 = v"1.10" <= VERSION < v"1.11"
                     is_1_11 = v"1.11" <= VERSION < v"1.12"
-                    is_svi_vnv = varinfo isa SimpleVarInfo{<:DynamicPPL.VarNamedVector}
-                    is_svi_od = varinfo isa SimpleVarInfo{<:OrderedDict}
+                    is_svi_vnv =
+                        linked_varinfo isa SimpleVarInfo{<:DynamicPPL.VarNamedVector}
+                    is_svi_od = linked_varinfo isa SimpleVarInfo{<:OrderedDict}
 
                     # Mooncake doesn't work with several combinations of SimpleVarInfo.
                     if is_mooncake && is_1_11 && is_svi_vnv
@@ -57,11 +58,10 @@ using DynamicPPL: LogDensityFunction
                             ref_ldf, adtype
                         )
                     else
-                        # TODO: Should we test linked as well?
                         @test DynamicPPL.TestUtils.AD.run_ad(
                             m,
                             adtype;
-                            varinfo=varinfo,
+                            varinfo=linked_varinfo,
                             expected_value_and_grad=(ref_logp, ref_grad),
                         ) isa Any
                     end