Merge branch 'main' into remove-logdensityfunctionwrapper-v2

Author: anurag-mds

HISTORY.md

@@ -5,6 +5,10 @@ Removed `getargnames`, `getmissings`, and `Base.nameof(::Model)` from the public
 Removed `LogDensityFunctionWrapper` and `VarInfo(::MarginalLogDensity, ...)` from the MarginalLogDensities extension.
 Users should now use `DynamicPPL.InitFromVector(mld, ...)` to obtain an initialisation strategy and pass it to `init!!` to get a consistent `VarInfo`.
 
+# 0.40.18
+
+Added a check on `unflatten!!` to error if the input vector was too long.
+
 # 0.40.17
 
 Implemented missing methods for `Base.copy` on internal structs.

src/varinfo.jl

@@ -537,6 +537,14 @@ end
 function unflatten!!(vi::VarInfo, vec::AbstractVector)
     vci = VectorChunkIterator!(vec, 1)
     new_values = map_values!!(vci, vi.values)
+    expected_len = vci.index - 1
+    if length(vec) != expected_len
+        throw(
+            DimensionMismatch(
+                "expected a vector of length $(expected_len), but got length $(length(vec))"
+            ),
+        )
+    end
     return VarInfo(vi.transform_strategy, new_values, vi.accs)
 end
 

test/logdensityfunction.jl

@@ -83,16 +83,30 @@ using Mooncake: Mooncake
     end
 end
 
-@testset "LogDensityFunction: Correctness" begin
+@testset "LogDensityFunction: correctness with multiple threads" begin
+    @testset "Threaded assume" begin
+        @model function threaded_assume()
+            x = zeros(10)
+            Threads.@threads for i in eachindex(x)
+                x[i] ~ Normal()
+            end
+        end
+        model = setthreadsafe(threaded_assume(), true)
+        ldf = DynamicPPL.LogDensityFunction(model)
+        xs = rand(ldf)
+        @test xs isa Vector{Float64} && length(xs) == 10
+        @test LogDensityProblems.logdensity(ldf, xs) ≈ sum(logpdf.(Normal(), xs))
+    end
+
     @testset "Threaded observe" begin
-        @model function threaded(y)
+        @model function threaded_observe(y)
             x ~ Normal()
             Threads.@threads for i in eachindex(y)
                 y[i] ~ Normal(x)
             end
         end
         N = 100
-        model = setthreadsafe(threaded(zeros(N)), true)
+        model = setthreadsafe(threaded_observe(zeros(N)), true)
         ldf = DynamicPPL.LogDensityFunction(model)
 
         xs = [1.0]

test/threadsafe.jl

@@ -111,22 +111,6 @@ const gdemo_default = gdemo_d()
         end
     end
 
-    @testset "promotion of VNT accumulators in TSVI" begin
-        # See https://github.com/TuringLang/DynamicPPL.jl/pull/1284.
-        @model function f()
-            x = zeros(10)
-            for i in eachindex(x)
-                x[i] ~ Normal()
-            end
-        end
-        model = setthreadsafe(f(), true)
-
-        vi = OnlyAccsVarInfo(RawValueAccumulator(false))
-        _, vi = DynamicPPL.init!!(model, vi, InitFromPrior(), UnlinkAll())
-        vnt = get_raw_values(vi)
-        @test vnt[@varname(x)] isa Vector{Float64}
-    end
-
     @testset "check_model with threadsafe" begin
         # This is a partial test for https://github.com/TuringLang/DynamicPPL.jl/issues/1157
         @model function f()
@@ -138,6 +122,34 @@ const gdemo_default = gdemo_d()
         @test !check_model(model)
     end
 
+    @testset "assumes are threadsafe" begin
+        # See https://github.com/TuringLang/DynamicPPL.jl/pull/1284.
+        #
+        # Note: anything that involves VarInfo is still thread-unsafe. But anything
+        # that uses OnlyAccsVarInfo is fine
+        @model function threaded_assume()
+            x = zeros(10)
+            Threads.@threads for i in eachindex(x)
+                x[i] ~ Normal()
+            end
+        end
+        model = setthreadsafe(threaded_assume(), true)
+
+        @testset "rand" begin
+            vnt = rand(model)
+            for i in 1:10
+                @test haskey(vnt, @varname(x[i]))
+            end
+        end
+        @testset "logprob" begin
+            xfixed = rand(10)
+            params = VarNamedTuple(; x=xfixed)
+            @test logprior(model, params) ≈ sum(logpdf.(Normal(), xfixed))
+            @test iszero(loglikelihood(model, params))
+            @test logjoint(model, params) ≈ sum(logpdf.(Normal(), xfixed))
+        end
+    end
+
     @testset "logprob correctness" begin
         x = rand(10_000)
 

test/varinfo.jl

@@ -115,7 +115,7 @@ end
         lp_d = logpdf(Normal(), values.d)
         m = demo() | (; c=values.c, d=values.d)
 
-        vi = DynamicPPL.unflatten!!(VarInfo(m), collect(values))
+        vi = DynamicPPL.unflatten!!(VarInfo(m), [values.a, values.b])
 
         vi = last(DynamicPPL.evaluate_nowarn!!(m, deepcopy(vi)))
         @test getlogprior(vi) == lp_a + lp_b
@@ -453,6 +453,21 @@ end
         end
     end
 
+    @testset "unflatten!! length check" begin
+        @model function demo_lc()
+            x ~ Normal()
+            return y ~ Normal(x, 1)
+        end
+        model = demo_lc() | (; y=0.0)
+        varinfo = VarInfo(model)
+        n = length(varinfo[:])
+        # Correct length should work.
+        @test DynamicPPL.unflatten!!(varinfo, zeros(n)) isa VarInfo
+        # Too many parameters should throw a DimensionMismatch.
+        @test_throws DimensionMismatch DynamicPPL.unflatten!!(varinfo, zeros(n + 1))
+        @test_throws DimensionMismatch DynamicPPL.unflatten!!(varinfo, zeros(2n))
+    end
+
     @testset "unflatten!! type stability" begin
         @model function demo(y)
             x ~ Normal()