Fix for tests in #516

src/test_utils.jl

@@ -51,9 +51,10 @@ function setup_varinfos(model::Model, example_values::NamedTuple, varnames)
     svi_typed = SimpleVarInfo(example_values)
     svi_untyped = SimpleVarInfo(OrderedDict())
 
+    lp = getlogp(vi_typed)
     return map((vi_untyped, vi_typed, svi_typed, svi_untyped)) do vi
         # Set them all to the same values.
-        update_values!!(vi, example_values, varnames)
+        DynamicPPL.setlogp!!(update_values!!(vi, example_values, varnames), lp)
     end
 end
 

test/linking.jl

@@ -91,21 +91,26 @@ end
         end
     end
 
+    # Related: https://github.com/TuringLang/DynamicPPL.jl/issues/504
     @testset "dirichlet" begin
-        @model demo_dirichlet() = x ~ Dirichlet(2, 1.0)
-        model = demo_dirichlet()
-        vis = DynamicPPL.TestUtils.setup_varinfos(model, rand(model), (@varname(x),))
-        @testset "$(short_varinfo_name(vi))" for vi in vis
-            @test length(vi[:]) == 2
-            @test iszero(getlogp(vi))
-            # Linked.
-            vi_linked = DynamicPPL.link!!(deepcopy(vi), model)
-            @test length(vi_linked[:]) == 1
-            @test !iszero(getlogp(vi_linked)) # should now include the log-absdet-jacobian correction
-            # Invlinked.
-            vi_invlinked = DynamicPPL.invlink!!(deepcopy(vi_linked), model)
-            @test length(vi_invlinked[:]) == 2
-            @test iszero(getlogp(vi_invlinked))
+        @model demo_dirichlet(d::Int) = x ~ Dirichlet(d, 1.0)
+        @testset "d=$d" for d in [2, 3, 5]
+            model = demo_dirichlet(d)
+            vis = DynamicPPL.TestUtils.setup_varinfos(model, rand(model), (@varname(x),))
+            @testset "$(short_varinfo_name(vi))" for vi in vis
+                lp = logpdf(Dirichlet(d, 1.0), vi[:])
+                @test length(vi[:]) == d
+                @test getlogp(vi) ≈ lp
+                # Linked.
+                vi_linked = DynamicPPL.link!!(deepcopy(vi), model)
+                @test length(vi_linked[:]) == d - 1
+                # Should now include the log-absdet-jacobian correction.
+                @test !(getlogp(vi_linked) ≈ lp)
+                # Invlinked.
+                vi_invlinked = DynamicPPL.invlink!!(deepcopy(vi_linked), model)
+                @test length(vi_invlinked[:]) == d
+                @test getlogp(vi_invlinked) ≈ lp
+            end
         end
     end
 end

test/varinfo.jl

@@ -273,40 +273,6 @@
         @test vals_prev == vi.metadata.x.vals
     end
 
-    # See https://github.com/TuringLang/DynamicPPL.jl/issues/504
-    @testset "Dimentionality checks" begin
-        @model function demo2d()
-            return x ~ Dirichlet(2, 1.0)
-        end
-        model = demo2d()
-        vi = VarInfo(model)                                   # make VarInfo -> sample from prior and compute logdensity
-        getlogp(vi) ≈ 0.0                                        # zero because Dirichlet(1) == Uniform over Simplex
-        spl = SampleFromPrior()                                # create dummy sampler for linking
-        DynamicPPL.link!!(vi, spl, model)                        # transform to unconstrained space
-        !(0.0 ≈ getlogp(last(DynamicPPL.evaluate!!(model, vi)))) # non-zero now due to log(abs(determinant(jacobian)))
-        x = vi[spl]                                              # extract unconstrained values
-        newx = deepcopy(x)                                      # simulate making a change to x
-        vinew = deepcopy(vi)
-        vinew[spl] = newx
-        @test vinew[spl] == newx
-
-        @model function demo3d()
-            return x ~ Dirichlet(3, 1.0)                                 # increase K to 3
-        end
-        model = demo3d()
-        vi = VarInfo(model)                                   # make VarInfo -> sample from prior and compute logdensity
-        getlogp(vi) ≈ 0.0                                        # zero because Dirichlet(1) == Uniform over Simplex
-        spl = SampleFromPrior()                                # create dummy sampler for linking
-        DynamicPPL.link!!(vi, spl, model)                        # transform to unconstrained space
-        !(0.0 ≈ getlogp(last(DynamicPPL.evaluate!!(model, vi)))) # non-zero now due to log(abs(determinant(jacobian)))
-        x = vi[spl]                                              # extract unconstrained values
-        newx = deepcopy(x)                                      # simulate making a change to x
-        vinew = deepcopy(vi)
-        vinew[spl] = newx
-
-        @test vinew[spl] == newx
-    end
-
     @testset "istrans" begin
         @model demo_constrained() = x ~ truncated(Normal(), 0, Inf)
         model = demo_constrained()