Remove SamplingContext from JETExt to avoid triggering Sampling->Init pathway

Author: penelopeysm

ext/DynamicPPLJETExt.jl

@@ -21,22 +21,17 @@ end
 function DynamicPPL.Experimental._determine_varinfo_jet(
     model::DynamicPPL.Model; only_ddpl::Bool=true
 )
-    # Use SamplingContext to test type stability.
-    sampling_model = DynamicPPL.contextualize(
-        model, DynamicPPL.SamplingContext(model.context)
-    )
-
     # First we try with the typed varinfo.
-    varinfo = DynamicPPL.typed_varinfo(sampling_model)
+    varinfo = DynamicPPL.typed_varinfo(model)
 
-    # Let's make sure that both evaluation and sampling doesn't result in type errors.
+    # Let's make sure that evaluation doesn't result in type errors.
     issuccess, result = DynamicPPL.Experimental.is_suitable_varinfo(
-        sampling_model, varinfo; only_ddpl
+        model, varinfo; only_ddpl
     )
 
     if !issuccess
         # Useful information for debugging.
-        @debug "Evaluaton with typed varinfo failed with the following issues:"
+        @debug "Evaluation with typed varinfo failed with the following issues:"
         @debug result
     end
 
@@ -46,7 +41,7 @@ function DynamicPPL.Experimental._determine_varinfo_jet(
     else
         # Warn the user that we can't use the type stable one.
         @warn "Model seems incompatible with typed varinfo. Falling back to untyped varinfo."
-        DynamicPPL.untyped_varinfo(sampling_model)
+        DynamicPPL.untyped_varinfo(model)
     end
 end
 

test/ext/DynamicPPLJETExt.jl

@@ -40,6 +40,11 @@
             end
         end
         @test DynamicPPL.Experimental.determine_suitable_varinfo(demo4()) isa
+            DynamicPPL.NTVarInfo
+        init_model = DynamicPPL.contextualize(
+            demo4(), DynamicPPL.InitContext(DynamicPPL.PriorInit())
+        )
+        @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.