Fix for #2205 (#2206)

* fix for #2205

* added tests

* bump patch version

* relax ADVI tests a bit

Project.toml

@@ -1,6 +1,6 @@
 name = "Turing"
 uuid = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
-version = "0.31.0"
+version = "0.31.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/variational/VariationalInference.jl

@@ -37,8 +37,11 @@ function make_logjoint(model::DynamicPPL.Model; weight = 1.0)
         DynamicPPL.DefaultContext(),
         weight
     )
-    model_contextualized = DynamicPPL.contextualize(model, ctx)
-    f = DynamicPPL.LogDensityFunction(model_contextualized)
+    f = DynamicPPL.LogDensityFunction(
+        model,
+        DynamicPPL.VarInfo(model),
+        ctx
+    )
     return Base.Fix1(LogDensityProblems.logdensity, f)
 end
 

test/variational/advi.jl

@@ -70,4 +70,24 @@
         x = rand(q, 1000)
         @test mean(eachcol(x)) ≈ [0.5, 0.5] atol=0.1
     end
+
+    # Ref: https://github.com/TuringLang/Turing.jl/issues/2205
+    @turing_testset "with `condition` (issue #2205)" begin
+        @model function demo_issue2205()
+            x ~ Normal()
+            y ~ Normal(x, 1)
+        end
+
+        model = demo_issue2205() | (y = 1.0,)
+        q = vi(model, ADVI(10, 1000))
+        # True mean.
+        mean_true = 1 / 2
+        var_true = 1 / 2
+        # Check the mean and variance of the posterior.
+        samples = rand(q, 1000)
+        mean_est = mean(samples)
+        var_est = var(samples)
+        @test mean_est ≈ mean_true atol=0.2
+        @test var_est ≈ var_true atol=0.2
+    end
 end