moved the JET.jl-dependent experimental determine_varinfo into a

separate `Experimental` module, as discussed

Author: torfjelde

ext/DynamicPPLJETExt.jl

@@ -3,7 +3,7 @@ module DynamicPPLJETExt
 using DynamicPPL: DynamicPPL
 using JET: JET
 
-function DynamicPPL.is_suitable_varinfo(
+function DynamicPPL.Experimental.is_suitable_varinfo(
     model::DynamicPPL.Model,
     context::DynamicPPL.AbstractContext,
     varinfo::DynamicPPL.AbstractVarInfo;
@@ -23,15 +23,17 @@ function DynamicPPL.is_suitable_varinfo(
     return length(JET.get_reports(result)) == 0, result
 end
 
-function DynamicPPL._determine_varinfo_jet(
+function DynamicPPL.Experimental._determine_varinfo_jet(
     model::DynamicPPL.Model, context::DynamicPPL.AbstractContext; only_ddpl::Bool=true
 )
     # First we try with the typed varinfo.
     varinfo = DynamicPPL.typed_varinfo(model, context)
     issuccess = true
 
     # Let's make sure that both evaluation and sampling doesn't result in type errors.
-    issuccess, result = DynamicPPL.is_suitable_varinfo(model, context, varinfo; only_ddpl)
+    issuccess, result = DynamicPPL.Experimental.is_suitable_varinfo(
+        model, context, varinfo; only_ddpl
+    )
 
     if !issuccess
         # Useful information for debugging.

src/DynamicPPL.jl

@@ -196,17 +196,19 @@ include("values_as_in_model.jl")
 include("debug_utils.jl")
 using .DebugUtils
 
+include("experimental.jl")
+
 # Better error message if users forget to load the AD package
 if isdefined(Base.Experimental, :register_error_hint)
     function __init__()
         Base.Experimental.register_error_hint(MethodError) do io, exc, argtypes, _
             requires_jet =
-                exc.f === _determine_varinfo_jet &&
+                exc.f === DynamicPPL.Experimental._determine_varinfo_jet &&
                 length(argtypes) >= 2 &&
                 argtypes[1] <: Model &&
                 argtypes[2] <: AbstractContext
             requires_jet |=
-                exc.f === is_suitable_varinfo &&
+                exc.f === DynamicPPL.Experimental.is_suitable_varinfo &&
                 length(argtypes) >= 3 &&
                 argtypes[1] <: Model &&
                 argtypes[2] <: AbstractContext &&

src/model_utils.jl

@@ -207,60 +207,3 @@ function value_iterator_from_chain(vi::AbstractVarInfo, chain)
         values_from_chain!(vi, chain, chain_idx, iteration_idx, OrderedDict{VarName,Any}())
     end
 end
-
-"""
-    is_suitable_varinfo(model::Model, context::AbstractContext, varinfo::AbstractVarInfo; kwargs...)
-
-Check if the `model` supports evaluation using the provided `context` and `varinfo`.
-
-!!! warning
-    Loading JET.jl is required before calling this function.
-
-# Arguments
-- `model`: The model to to verify the support for.
-- `context`: The context to use for the model evaluation.
-- `varinfo`: The varinfo to verify the support for.
-
-# Keyword Arguments
-- `only_ddpl`: If `true`, only consider error reports occuring in the tilde pipeline. Default: `true`.
-
-# Returns
-- `issuccess`: `true` if the model supports the varinfo, otherwise `false`.
-- `report`: The result of `report_call` from JET.jl.
-"""
-function is_suitable_varinfo end
-
-# Internal hook for JET.jl to overload.
-function _determine_varinfo_jet end
-
-"""
-    determine_suitable_varinfo(model[, context]; verbose::Bool=false, only_ddpl::Bool=true)
-
-Return a suitable varinfo for the given `model`.
-
-See also: [`DynamicPPL.is_suitable_varinfo`](@ref).
-
-!!! warning
-    For full functionality, this requires JET.jl to be loaded.
-    If JET.jl is not loaded, this function will assume the model is compatible with typed varinfo.
-
-# Arguments
-- `model`: The model for which to determine the varinfo.
-- `context`: The context to use for the model evaluation. Default: `SamplingContext()`.
-
-# Keyword Arguments
-- `only_ddpl`: If `true`, only consider error reports within DynamicPPL.jl.
-"""
-function determine_suitable_varinfo(
-    model::Model, context::AbstractContext=SamplingContext(); only_ddpl::Bool=true
-)
-    # If JET.jl has been loaded, and thus `determine_varinfo` has been defined, we use that.
-    return if Base.get_extension(DynamicPPL, :DynamicPPLJETExt) !== nothing
-        _determine_varinfo_jet(model, context; only_ddpl)
-    else
-        # Warn the user.
-        @warn "JET.jl is not loaded. Assumes the model is compatible with typed varinfo."
-        # Otherwise, we use the, possibly incorrect, default typed varinfo (to stay backwards compat).
-        typed_varinfo(model, context)
-    end
-end

test/ext/DynamicPPLJETExt.jl

@@ -9,10 +9,12 @@
             end
         end
         model = demo1()
-        @test DynamicPPL.determine_suitable_varinfo(model) isa DynamicPPL.UntypedVarInfo
+        @test DynamicPPL.Experimental.determine_suitable_varinfo(model) isa
+            DynamicPPL.UntypedVarInfo
 
         @model demo2() = x ~ Normal()
-        @test DynamicPPL.determine_suitable_varinfo(demo2()) isa DynamicPPL.TypedVarInfo
+        @test DynamicPPL.Experimental.determine_suitable_varinfo(demo2()) isa
+            DynamicPPL.TypedVarInfo
 
         @model function demo3()
             # Just making sure that nothing strange happens when type inference fails.
@@ -24,7 +26,8 @@
                 z ~ Normal()
             end
         end
-        @test DynamicPPL.determine_suitable_varinfo(demo3()) isa DynamicPPL.UntypedVarInfo
+        @test DynamicPPL.Experimental.determine_suitable_varinfo(demo3()) isa
+            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}}`.
@@ -36,7 +39,8 @@
                 y ~ Cauchy() # different distibution, but same transformation
             end
         end
-        @test DynamicPPL.determine_suitable_varinfo(demo4()) isa DynamicPPL.UntypedVarInfo
+        @test DynamicPPL.Experimental.determine_suitable_varinfo(demo4()) isa
+            DynamicPPL.UntypedVarInfo
 
         # In this model, the type error occurs in the user code rather than in DynamicPPL.
         @model function demo5()
@@ -48,16 +52,18 @@
             return sum(xs)
         end
         # Should pass if we're only checking the tilde statements.
-        @test DynamicPPL.determine_suitable_varinfo(demo5()) isa DynamicPPL.TypedVarInfo
+        @test DynamicPPL.Experimental.determine_suitable_varinfo(demo5()) isa
+            DynamicPPL.TypedVarInfo
         # Should fail if we're including errors in the model body.
-        @test DynamicPPL.determine_suitable_varinfo(demo5(); only_ddpl=false) isa
-            DynamicPPL.UntypedVarInfo
+        @test DynamicPPL.Experimental.determine_suitable_varinfo(
+            demo5(); only_ddpl=false
+        ) isa DynamicPPL.UntypedVarInfo
     end
 
     @testset "demo models" begin
         @testset "$(model.f)" for model in DynamicPPL.TestUtils.DEMO_MODELS
             # Use debug logging below.
-            varinfo = DynamicPPL.DynamicPPL.determine_suitable_varinfo(model)
+            varinfo = DynamicPPL.Experimental.determine_suitable_varinfo(model)
             # They should all result in typed.
             @test varinfo isa DynamicPPL.TypedVarInfo
             # But let's also make sure that they're not lying.