determine_suitable_varinfo now only performs checks using the

provided context, but uses `SamplingContext` by default (as this
should be a stricter check than just evaluation)

Author: torfjelde

ext/DynamicPPLJETExt.jl

@@ -50,37 +50,29 @@ end
 
 function DynamicPPL._determine_varinfo_jet(
     model::DynamicPPL.Model,
-    context::DynamicPPL.AbstractContext=DynamicPPL.DefaultContext();
-    verbose::Bool=false,
+    context::DynamicPPL.AbstractContext;
     only_tilde::Bool=true,
 )
     # First we try with the typed varinfo.
-    varinfo = DynamicPPL.typed_varinfo(model)
+    varinfo = if DynamicPPL.hassampler(context)
+        # Don't need to add sampling context for this to work.
+        DynamicPPL.typed_varinfo(model, context)
+    else
+        # Need a sampling context to initialize the varinfo.
+        DynamicPPL.typed_varinfo(model, DynamicPPL.SamplingContext(context))
+    end
     issuccess = true
 
     # Let's make sure that both evaluation and sampling doesn't result in type errors.
-    issuccess, reports_eval = DynamicPPL.is_suitable_varinfo(
+    issuccess, reports = DynamicPPL.is_suitable_varinfo(
         model, context, varinfo; only_tilde
     )
 
-    if issuccess
-        # Evaluation succeeded, let's try sampling.
-        issuccess_sample, reports_sample = DynamicPPL.is_suitable_varinfo(
-            model, DynamicPPL.SamplingContext(context), varinfo; only_tilde
-        )
-        issuccess &= issuccess_sample
-        if !issuccess && verbose
-            # Show the user the issues.
-            @warn "Sampling with typed varinfo failed with the following issues:"
-            for report in reports_sample
-                @warn report
-            end
-        end
-    elseif verbose
-        # Show the user the issues.
-        @warn "Evaluaton with typed varinfo failed with the following issues:"
-        for report in reports_eval
-            @warn report
+    if !issuccess
+        # Useful information for debugging.
+        @debug "Evaluaton with typed varinfo failed with the following issues:"
+        for report in reports
+            @debug report
         end
     end
 

src/model_utils.jl

@@ -246,24 +246,19 @@ See also: [`DynamicPPL.is_suitable_varinfo`](@ref).
 
 # Arguments
 - `model`: The model for which to determine the varinfo.
-- `context`: The context to use for the evaluation and sampling. Default: `DefaultContext()`.
+- `context`: The context to use for the model evaluation. Default: `SamplingContext()`.
 
 # Keyword Arguments
-- `verbose`: If `true`, the user will be warned about the issues that caused the fallback to untyped varinfo.
 - `only_tilde`: If `true`, only consider error reports occuring in the tilde pipeline. Default: `true`.
-
-# Keyword Arguments
-- `verbose`: If `true`, the user will be warned about the issues that caused the fallback to untyped varinfo.
 """
 function determine_suitable_varinfo(
     model::Model,
-    context::AbstractContext=DefaultContext();
-    verbose::Bool=false,
+    context::AbstractContext=SamplingContext();
     only_tilde::Bool=true,
 )
     # If JET.jl has been loaded, and thus `determine_varinfo` has been defined, we use that.
     if Base.get_extension(DynamicPPL, :DynamicPPLJETExt) !== nothing
-        return _determine_varinfo_jet(model, context; only_tilde, verbose)
+        return _determine_varinfo_jet(model, context; only_tilde)
     else
         @warn "JET.jl is not loaded. Assumes the model is compatible with typed varinfo."
     end

src/sampler.jl

@@ -77,33 +77,7 @@ function default_varinfo(
     context::AbstractContext,
 )
     init_sampler = initialsampler(sampler)
-    varinfo = VarInfo(rng, model, init_sampler, context)
-
-    # If JET.jl has been loaded => use static checking to see if we can actually use the typed varinfo.
-    if Base.get_extension(DynamicPPL, :DynamicPPLJETExt) !== nothing
-        # Check evaluation.
-        issuccess = first(is_suitable_varinfo(model, context, varinfo))
-        if issuccess
-            # Check the initial sampler.
-            issuccess &= first(
-                is_suitable_varinfo(model, SamplingContext(init_sampler, context), varinfo)
-            )
-        end
-
-        if issuccess
-            # Check the actual sampler.
-            issuccess &= first(
-                is_suitable_varinfo(model, SamplingContext(sampler, context), varinfo)
-            )
-        end
-
-        if !issuccess
-            @warn "Model seems incompatible with typed varinfo. Falling back to untyped varinfo."
-            # TODO: Use a constructor which takes the rng and the sampler too.
-            varinfo = untyped_varinfo(model)
-        end
-    end
-
+    varinfo = determine_suitable_varinfo(model, SamplingContext(rng, init_sampler, context))
     return varinfo
 end