Fix JETExt properly

Author: penelopeysm

ext/DynamicPPLJETExt.jl

@@ -6,7 +6,6 @@ using JET: JET
 function DynamicPPL.Experimental.is_suitable_varinfo(
     model::DynamicPPL.Model, varinfo::DynamicPPL.AbstractVarInfo; only_ddpl::Bool=true
 )
-    # Let's make sure that both evaluation and sampling doesn't result in type errors.
     f, argtypes = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(model, varinfo)
     # If specified, we only check errors originating somewhere in the DynamicPPL.jl.
     # This way we don't just fall back to untyped if the user's code is the issue.

src/context_implementations.jl

@@ -28,13 +28,6 @@ end
 function tilde_assume(rng::Random.AbstractRNG, ::DefaultContext, sampler, right, vn, vi)
     return assume(rng, sampler, right, vn, vi)
 end
-function tilde_assume(rng::Random.AbstractRNG, ::InitContext, sampler, right, vn, vi)
-    @warn(
-        "Encountered SamplingContext->InitContext. This method will be removed in the next PR.",
-    )
-    # just pretend the `InitContext` isn't there for now.
-    return assume(rng, sampler, right, vn, vi)
-end
 function tilde_assume(::DefaultContext, sampler, right, vn, vi)
     # same as above but no rng
     return assume(Random.default_rng(), sampler, right, vn, vi)

test/ext/DynamicPPLJETExt.jl

@@ -30,7 +30,7 @@
             DynamicPPL.UntypedVarInfo
 
         # Evaluation works (and it would even do so in practice), but sampling
-        # fill fail due to storing `Cauchy{Float64}` in `Vector{Normal{Float64}}`.
+        # will fail due to storing `Cauchy{Float64}` in `Vector{Normal{Float64}}`.
         @model function demo4()
             x ~ Bernoulli()
             if x
@@ -40,11 +40,6 @@
             end
         end
         @test DynamicPPL.Experimental.determine_suitable_varinfo(demo4()) isa
-            DynamicPPL.NTVarInfo
-        init_model = DynamicPPL.contextualize(
-            demo4(), DynamicPPL.InitContext(DynamicPPL.InitFromPrior())
-        )
-        @test DynamicPPL.Experimental.determine_suitable_varinfo(init_model) isa
             DynamicPPL.UntypedVarInfo
 
         # In this model, the type error occurs in the user code rather than in DynamicPPL.
@@ -67,33 +62,37 @@
 
     @testset "demo models" begin
         @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
-            sampling_model = contextualize(model, SamplingContext(model.context))
             # Use debug logging below.
             varinfo = DynamicPPL.Experimental.determine_suitable_varinfo(model)
-            # Check that the inferred varinfo is indeed suitable for evaluation and sampling
-            f_eval, argtypes_eval = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                model, varinfo
-            )
-            JET.test_call(f_eval, argtypes_eval)
-
-            f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                sampling_model, varinfo
-            )
-            JET.test_call(f_sample, argtypes_sample)
             # For our demo models, they should all result in typed.
             is_typed = varinfo isa DynamicPPL.NTVarInfo
             @test is_typed
-            # If the test failed, check why it didn't infer a typed varinfo
+            # If the test failed, check what the type stability problem was for
+            # the typed varinfo. This is mostly useful for debugging from test
+            # logs.
             if !is_typed
+                @info "Model `$(model.f)` is not type stable with typed varinfo."
                 typed_vi = DynamicPPL.typed_varinfo(model)
-                f_eval, argtypes_eval = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                    model, typed_vi
+
+                @info "Evaluating with DefaultContext:"
+                model = DynamicPPL.contextualize(
+                    model,
+                    DynamicPPL.setleafcontext(model.context, DynamicPPL.DefaultContext()),
+                )
+                f, argtypes = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
+                    model, varinfo
+                )
+                JET.test_call(f, argtypes)
+
+                @info "Initialising with InitContext:"
+                model = DynamicPPL.contextualize(
+                    model,
+                    DynamicPPL.setleafcontext(model.context, DynamicPPL.InitContext()),
                 )
-                JET.test_call(f_eval, argtypes_eval)
-                f_sample, argtypes_sample = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
-                    sampling_model, typed_vi
+                f, argtypes = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
+                    model, varinfo
                 )
-                JET.test_call(f_sample, argtypes_sample)
+                JET.test_call(f, argtypes)
             end
         end
     end