Add more tests for dot tilde implementations (#247)

This PR extends the tests that were introduced in #245.

Co-authored-by: David Widmann <[email protected]>

Author: devmotion

test/Project.toml

@@ -7,6 +7,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DistributionsAD = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

test/context_implementations.jl

@@ -15,26 +15,58 @@
     end
 
     # https://github.com/TuringLang/DynamicPPL.jl/issues/28#issuecomment-829223577
-    @testset "arrays of distributions" begin
-        @model function test(x, y)
-            y .~ Normal.(x)
+    @testset "dot tilde: arrays of distributions" begin
+        @testset "assume" begin
+            @model function test(x, size)
+                y = Array{Float64,length(size)}(undef, size...)
+                y .~ Normal.(x)
+                return y, getlogp(__varinfo__)
+            end
+
+            for ysize in ((2,), (2, 3), (2, 3, 4))
+                for x in (
+                    # scalar
+                    randn(),
+                    # drop trailing dimensions
+                    ntuple(i -> randn(ysize[1:i]), length(ysize))...,
+                    # singleton dimensions
+                    ntuple(
+                        i -> randn(ysize[1:(i-1)]..., 1, ysize[(i+1):end]...),
+                        length(ysize),
+                    )...,
+                )
+                    model = test(x, ysize)
+                    y, lp = model()
+                    @test lp ≈ sum(logpdf.(Normal.(x), y))
+
+                    ys = [first(model()) for _ in 1:10_000]
+                    @test norm(mean(ys) .- x, Inf) < 0.1
+                    @test norm(std(ys) .- 1, Inf) < 0.1
+                end
+            end
         end
 
-        for ysize in ((2,), (2, 3), (2, 3, 4))
-            # drop trailing dimensions
-            for xsize in ntuple(i -> ysize[1:i], length(ysize))
-                x = randn(xsize)
-                y = randn(ysize)
-                z = logjoint(test(x, y), VarInfo())
-                @test z ≈ sum(logpdf.(Normal.(x), y))
+        @testset "observe" begin
+            @model function test(x, y)
+                y .~ Normal.(x)
             end
 
-            # singleton dimensions
-            for xsize in ntuple(i -> (ysize[1:(i-1)]..., 1, ysize[(i+1):end]...), length(ysize))
-                x = randn(xsize)
-                y = randn(ysize)
-                z = logjoint(test(x, y), VarInfo())
-                @test z ≈ sum(logpdf.(Normal.(x), y))
+            for ysize in ((2,), (2, 3), (2, 3, 4))
+                for x in (
+                    # scalar
+                    randn(),
+                    # drop trailing dimensions
+                    ntuple(i -> randn(ysize[1:i]), length(ysize))...,
+                    # singleton dimensions
+                    ntuple(
+                        i -> randn(ysize[1:(i-1)]..., 1, ysize[(i+1):end]...),
+                        length(ysize),
+                    )...,
+                )
+                    y = randn(ysize)
+                    z = logjoint(test(x, y), VarInfo())
+                    @test z ≈ sum(logpdf.(Normal.(x), y))
+                end
             end
         end
     end

test/runtests.jl

@@ -12,6 +12,7 @@ using Tracker
 using Zygote
 
 using Distributed
+using LinearAlgebra
 using Pkg
 using Random
 using Serialization