final fixes?

Author: penelopeysm

src/debug_utils.jl

@@ -131,9 +131,7 @@ A context used for checking validity of a model.
 # Fields
 $(FIELDS)
 """
-struct DebugContext{M<:Model,C<:AbstractContext} <: AbstractContext
-    "model that is being run"
-    model::M
+struct DebugContext{C<:AbstractContext} <: AbstractContext
     "context used for running the model"
     context::C
     "mapping from varnames to the number of times they have been seen"
@@ -149,7 +147,6 @@ struct DebugContext{M<:Model,C<:AbstractContext} <: AbstractContext
 end
 
 function DebugContext(
-    model::Model,
     context::AbstractContext=DefaultContext();
     varnames_seen=OrderedDict{VarName,Int}(),
     statements=Vector{Stmt}(),
@@ -158,7 +155,6 @@ function DebugContext(
     record_varinfo=false,
 )
     return DebugContext(
-        model,
         context,
         varnames_seen,
         statements,
@@ -344,7 +340,7 @@ function check_varnames_seen(varnames_seen::AbstractDict{VarName,Int})
 end
 
 # A check we run on the model before evaluating it.
-function check_model_pre_evaluation(context::DebugContext, model::Model)
+function check_model_pre_evaluation(model::Model)
     issuccess = true
     # If something is in the model arguments, then it should NOT be in `condition`,
     # nor should there be any symbol present in `condition` that has the same symbol.
@@ -361,8 +357,8 @@ function check_model_pre_evaluation(context::DebugContext, model::Model)
     return issuccess
 end
 
-function check_model_post_evaluation(context::DebugContext, model::Model)
-    return check_varnames_seen(context.varnames_seen)
+function check_model_post_evaluation(model::Model)
+    return check_varnames_seen(model.context.varnames_seen)
 end
 
 """
@@ -443,21 +439,18 @@ function check_model_and_trace(
 )
     # Execute the model with the debug context.
     debug_context = DebugContext(
-        model,
-        SamplingContext(rng, model.context);
-        error_on_failure=error_on_failure,
-        kwargs...,
+        SamplingContext(rng, model.context); error_on_failure=error_on_failure, kwargs...
     )
     debug_model = DynamicPPL.contextualize(model, debug_context)
 
     # Perform checks before evaluating the model.
-    issuccess = check_model_pre_evaluation(debug_context, debug_model)
+    issuccess = check_model_pre_evaluation(debug_model)
 
     # Force single-threaded execution.
     DynamicPPL.evaluate_threadunsafe!!(debug_model, varinfo)
 
     # Perform checks after evaluating the model.
-    issuccess &= check_model_post_evaluation(debug_context, debug_model)
+    issuccess &= check_model_post_evaluation(debug_model)
 
     if !issuccess && error_on_failure
         error("model check failed")

test/compiler.jl

@@ -193,7 +193,10 @@ module Issue537 end
         varinfo = VarInfo(model)
         @test getlogjoint(varinfo) == lp
         @test varinfo_ isa AbstractVarInfo
-        @test model_ === model
+        # During the model evaluation, its context is wrapped in a
+        # SamplingContext, so `model_` is not going to be equal to `model`.
+        # We can still check equality of `f` though.
+        @test model_.f === model.f
         @test model_.context isa SamplingContext
         @test model_.context.rng isa Random.AbstractRNG
 

test/debug_utils.jl

@@ -1,7 +1,7 @@
 @testset "check_model" begin
     @testset "context interface" begin
         @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
-            context = DynamicPPL.DebugUtils.DebugContext(model)
+            context = DynamicPPL.DebugUtils.DebugContext()
             DynamicPPL.TestUtils.test_context(context, model)
         end
     end
@@ -35,9 +35,7 @@
             buggy_model = buggy_demo_model()
 
             @test_logs (:warn,) (:warn,) check_model(buggy_model)
-            issuccess = check_model(
-                buggy_model; context=SamplingContext(), record_varinfo=false
-            )
+            issuccess = check_model(buggy_model; record_varinfo=false)
             @test !issuccess
             @test_throws ErrorException check_model(buggy_model; error_on_failure=true)
         end
@@ -81,9 +79,7 @@
             buggy_model = buggy_subsumes_demo_model()
 
             @test_logs (:warn,) (:warn,) check_model(buggy_model)
-            issuccess = check_model(
-                buggy_model; context=SamplingContext(), record_varinfo=false
-            )
+            issuccess = check_model(buggy_model; record_varinfo=false)
             @test !issuccess
             @test_throws ErrorException check_model(buggy_model; error_on_failure=true)
         end
@@ -98,9 +94,7 @@
             buggy_model = buggy_subsumes_demo_model()
 
             @test_logs (:warn,) (:warn,) check_model(buggy_model)
-            issuccess = check_model(
-                buggy_model; context=SamplingContext(), record_varinfo=false
-            )
+            issuccess = check_model(buggy_model; record_varinfo=false)
             @test !issuccess
             @test_throws ErrorException check_model(buggy_model; error_on_failure=true)
         end
@@ -115,9 +109,7 @@
             buggy_model = buggy_subsumes_demo_model()
 
             @test_logs (:warn,) (:warn,) check_model(buggy_model)
-            issuccess = check_model(
-                buggy_model; context=SamplingContext(), record_varinfo=false
-            )
+            issuccess = check_model(buggy_model; record_varinfo=false)
             @test !issuccess
             @test_throws ErrorException check_model(buggy_model; error_on_failure=true)
         end

test/model.jl

@@ -591,7 +591,7 @@ const GDEMO_DEFAULT = DynamicPPL.TestUtils.demo_assume_observe_literal()
             xs_train = 1:0.1:10
             ys_train = ground_truth_β .* xs_train + rand(Normal(0, 0.1), length(xs_train))
             m_lin_reg = linear_reg(xs_train, ys_train)
-            chain = [last(evaluate!!(m_lin_reg, VarInfo())) for _ in 1:10000]
+            chain = [last(sample!!(m_lin_reg, VarInfo())) for _ in 1:10000]
 
             # chain is generated from the prior
             @test mean([chain[i][@varname(β)] for i in eachindex(chain)]) ≈ 1.0 atol = 0.1

test/threadsafe.jl

@@ -52,9 +52,10 @@
                 x[i] ~ Normal(x[i - 1], 1)
             end
         end
+        model = wthreads(x)
 
         vi = VarInfo()
-        wthreads(x)(vi)
+        model(vi)
         lp_w_threads = getlogjoint(vi)
         if Threads.nthreads() == 1
             @test vi_ isa VarInfo
@@ -64,23 +65,19 @@
 
         println("With `@threads`:")
         println("  default:")
-        @time wthreads(x)(vi)
+        @time model(vi)
 
         # Ensure that we use `ThreadSafeVarInfo` to handle multithreaded observe statements.
-        DynamicPPL.evaluate_threadsafe!!(
-            wthreads(x),
-            vi,
-            SamplingContext(Random.default_rng(), SampleFromPrior(), DefaultContext()),
-        )
+        sampling_model = contextualize(model, SamplingContext(model.context))
+        DynamicPPL.evaluate_threadsafe!!(sampling_model, vi)
         @test getlogjoint(vi) ≈ lp_w_threads
+        # check that it's wrapped during the model evaluation
         @test vi_ isa DynamicPPL.ThreadSafeVarInfo
+        # ensure that it's unwrapped after evaluation finishes
+        @test vi isa VarInfo
 
         println("  evaluate_threadsafe!!:")
-        @time DynamicPPL.evaluate_threadsafe!!(
-            wthreads(x),
-            vi,
-            SamplingContext(Random.default_rng(), SampleFromPrior(), DefaultContext()),
-        )
+        @time DynamicPPL.evaluate_threadsafe!!(sampling_model, vi)
 
         @model function wothreads(x)
             global vi_ = __varinfo__
@@ -89,9 +86,10 @@
                 x[i] ~ Normal(x[i - 1], 1)
             end
         end
+        model = wothreads(x)
 
         vi = VarInfo()
-        wothreads(x)(vi)
+        model(vi)
         lp_wo_threads = getlogjoint(vi)
         if Threads.nthreads() == 1
             @test vi_ isa VarInfo
@@ -101,24 +99,18 @@
 
         println("Without `@threads`:")
         println("  default:")
-        @time wothreads(x)(vi)
+        @time model(vi)
 
         @test lp_w_threads ≈ lp_wo_threads
 
         # Ensure that we use `VarInfo`.
-        DynamicPPL.evaluate_threadunsafe!!(
-            wothreads(x),
-            vi,
-            SamplingContext(Random.default_rng(), SampleFromPrior(), DefaultContext()),
-        )
+        sampling_model = contextualize(model, SamplingContext(model.context))
+        DynamicPPL.evaluate_threadunsafe!!(sampling_model, vi)
         @test getlogjoint(vi) ≈ lp_w_threads
         @test vi_ isa VarInfo
+        @test vi isa VarInfo
 
         println("  evaluate_threadunsafe!!:")
-        @time DynamicPPL.evaluate_threadunsafe!!(
-            wothreads(x),
-            vi,
-            SamplingContext(Random.default_rng(), SampleFromPrior(), DefaultContext()),
-        )
+        @time DynamicPPL.evaluate_threadunsafe!!(sampling_model, vi)
     end
 end