Fix JETExt

Author: penelopeysm

ext/DynamicPPLJETExt.jl

@@ -21,22 +21,35 @@ end
 function DynamicPPL.Experimental._determine_varinfo_jet(
     model::DynamicPPL.Model; only_ddpl::Bool=true
 )
-    # First we try with the typed varinfo.
+    # Generate a typed varinfo to test model type stability with
     varinfo = DynamicPPL.typed_varinfo(model)
 
-    # Let's make sure that evaluation doesn't result in type errors.
-    issuccess, result = DynamicPPL.Experimental.is_suitable_varinfo(
+    # Check type stability of evaluation (i.e. DefaultContext)
+    model = DynamicPPL.contextualize(
+        model, DynamicPPL.setleafcontext(model.context, DynamicPPL.DefaultContext())
+    )
+    eval_issuccess, eval_result = DynamicPPL.Experimental.is_suitable_varinfo(
         model, varinfo; only_ddpl
     )
-
-    if !issuccess
-        # Useful information for debugging.
+    if !eval_issuccess
         @debug "Evaluation with typed varinfo failed with the following issues:"
-        @debug result
+        @debug eval_result
+    end
+
+    # Check type stability of initialisation (i.e. InitContext)
+    model = DynamicPPL.contextualize(
+        model, DynamicPPL.setleafcontext(model.context, DynamicPPL.InitContext())
+    )
+    init_issuccess, init_result = DynamicPPL.Experimental.is_suitable_varinfo(
+        model, varinfo; only_ddpl
+    )
+    if !init_issuccess
+        @debug "Initialisation with typed varinfo failed with the following issues:"
+        @debug init_result
     end
 
-    # If we didn't fail anywhere, we return the type stable one.
-    return if issuccess
+    # If neither of them failed, we can return the typed varinfo as it's type stable.
+    return if (eval_issuccess && init_issuccess)
         varinfo
     else
         # Warn the user that we can't use the type stable one.