Using JET.jl to determine if typed varinfo is okay

Project.toml

@@ -15,6 +15,7 @@ ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
@@ -29,6 +30,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
@@ -37,6 +39,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 DynamicPPLChainRulesCoreExt = ["ChainRulesCore"]
 DynamicPPLEnzymeCoreExt = ["EnzymeCore"]
 DynamicPPLForwardDiffExt = ["ForwardDiff"]
+DynamicPPLJETExt = ["JET"]
 DynamicPPLMCMCChainsExt = ["MCMCChains"]
 DynamicPPLReverseDiffExt = ["ReverseDiff"]
 DynamicPPLZygoteRulesExt = ["ZygoteRules"]
@@ -55,6 +58,7 @@ Distributions = "0.25"
 DocStringExtensions = "0.9"
 EnzymeCore = "0.6 - 0.8"
 ForwardDiff = "0.10"
+JET = "0.9"
 LinearAlgebra = "1.6"
 LogDensityProblems = "2"
 LogDensityProblemsAD = "1.7.0"
@@ -72,6 +76,7 @@ julia = "1.10"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"

docs/Project.toml

@@ -6,6 +6,7 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterMermaid = "a078cd44-4d9c-4618-b545-3ab9d77f9177"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -18,6 +19,7 @@ Documenter = "1"
 DocumenterMermaid = "0.1"
 FillArrays = "0.13, 1"
 ForwardDiff = "0.10"
+JET = "0.9"
 LogDensityProblems = "2"
 MCMCChains = "5, 6"
 StableRNGs = "1"

docs/src/api.md

@@ -265,6 +265,13 @@ AbstractVarInfo
 
 But exactly how a [`AbstractVarInfo`](@ref) stores this information can vary.
 
+For constructing the "default" typed and untyped varinfo types used in DynamicPPL (see [the section on varinfo design](@ref "Design of `VarInfo`") for more on this), we have the following two methods:
+
+```@docs
+DynamicPPL.untyped_varinfo
+DynamicPPL.typed_varinfo
+```
+
 #### `VarInfo`
 
 ```@docs
@@ -425,6 +432,19 @@ DynamicPPL.loadstate
 DynamicPPL.initialsampler
 ```
 
+Finally, to specify which varinfo type a [`Sampler`](@ref) should use for a given [`Model`](@ref), this is specified by [`DynamicPPL.default_varinfo`](@ref) and can thus be overloaded for each  `model`-`sampler` combination. This can be useful in cases where one has explicit knowledge that one type of varinfo will be more performant for the given `model` and `sampler`.
+
+```@docs
+DynamicPPL.default_varinfo
+```
+
+There is also the _experimental_ [`DynamicPPL.Experimental.determine_suitable_varinfo`](@ref), which uses static checking via [JET.jl](https://github.com/aviatesk/JET.jl) to determine whether one should use [`DynamicPPL.typed_varinfo`](@ref) or [`DynamicPPL.untyped_varinfo`](@ref), depending on which supports the model:
+
+```@docs
+DynamicPPL.Experimental.determine_suitable_varinfo
+DynamicPPL.Experimental.is_suitable_varinfo
+```
+
 ### [Model-Internal Functions](@id model_internal)
 
 ```@docs

docs/src/internals/varinfo.md

@@ -79,9 +79,7 @@ For example, with the model above we have
 
 ```@example varinfo-design
 # Type-unstable `VarInfo`
-varinfo_untyped = DynamicPPL.untyped_varinfo(
-    demo(), SampleFromPrior(), DefaultContext(), DynamicPPL.Metadata()
-)
+varinfo_untyped = DynamicPPL.untyped_varinfo(demo())
 typeof(varinfo_untyped.metadata)
 ```
 

ext/DynamicPPLJETExt.jl

@@ -0,0 +1,54 @@
+module DynamicPPLJETExt
+
+using DynamicPPL: DynamicPPL
+using JET: JET
+
+function DynamicPPL.Experimental.is_suitable_varinfo(
+    model::DynamicPPL.Model,
+    context::DynamicPPL.AbstractContext,
+    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, context
+    )
+    # 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.
+    result = if only_ddpl
+        JET.report_call(f, argtypes; target_modules=(JET.AnyFrameModule(DynamicPPL),))
+    else
+        JET.report_call(f, argtypes)
+    end
+    return length(JET.get_reports(result)) == 0, result
+end
+
+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.Experimental.is_suitable_varinfo(
+        model, context, varinfo; only_ddpl
+    )
+
+    if !issuccess
+        # Useful information for debugging.
+        @debug "Evaluaton with typed varinfo failed with the following issues:"
+        @debug result
+    end
+
+    # If we didn't fail anywhere, we return the type stable one.
+    return if issuccess
+        varinfo
+    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(model, context)
+    end
+end
+
+end

src/DynamicPPL.jl

@@ -199,32 +199,35 @@ include("values_as_in_model.jl")
 include("debug_utils.jl")
 using .DebugUtils
 
+include("experimental.jl")
 include("deprecated.jl")
 
 if !isdefined(Base, :get_extension)
     using Requires
 end
 
-@static if !isdefined(Base, :get_extension)
+# Better error message if users forget to load the AD package
+if isdefined(Base.Experimental, :register_error_hint)
     function __init__()
-        @require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" include(
-            "../ext/DynamicPPLChainRulesCoreExt.jl"
-        )
-        @require EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" include(
-            "../ext/DynamicPPLEnzymeCoreExt.jl"
-        )
-        @require ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" include(
-            "../ext/DynamicPPLForwardDiffExt.jl"
-        )
-        @require MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d" include(
-            "../ext/DynamicPPLMCMCChainsExt.jl"
-        )
-        @require ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" include(
-            "../ext/DynamicPPLReverseDiffExt.jl"
-        )
-        @require ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444" include(
-            "../ext/DynamicPPLZygoteRulesExt.jl"
-        )
+        Base.Experimental.register_error_hint(MethodError) do io, exc, argtypes, _
+            requires_jet =
+                exc.f === DynamicPPL.Experimental._determine_varinfo_jet &&
+                length(argtypes) >= 2 &&
+                argtypes[1] <: Model &&
+                argtypes[2] <: AbstractContext
+            requires_jet |=
+                exc.f === DynamicPPL.Experimental.is_suitable_varinfo &&
+                length(argtypes) >= 3 &&
+                argtypes[1] <: Model &&
+                argtypes[2] <: AbstractContext &&
+                argtypes[3] <: AbstractVarInfo
+            if requires_jet
+                print(
+                    io,
+                    "\n$(exc.f) requires JET.jl to be loaded. Please run `using JET` before calling $(exc.f).",
+                )
+            end
+        end
     end
 end
 

src/experimental.jl

@@ -0,0 +1,103 @@
+module Experimental
+
+using DynamicPPL: DynamicPPL
+
+"""
+    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.Experimental.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.
+
+# Examples
+
+```jldoctest
+julia> using DynamicPPL.Experimental: determine_suitable_varinfo
+
+julia> using JET: JET  # needs to be loaded for full functionality
+
+julia> @model function model_with_random_support()
+           x ~ Bernoulli()
+           if x
+               y ~ Normal()
+           else
+               z ~ Normal()
+           end
+       end
+model_with_random_support (generic function with 2 methods)
+
+julia> model = model_with_random_support();
+
+julia> # Typed varinfo cannot handle this random support model properly
+       # as using a single execution of the model will not see all random variables.
+       # Hence, this this model requires untyped varinfo.
+       varinfo = determine_suitable_varinfo(model);
+┌ Warning: Model seems incompatible with typed varinfo. Falling back to untyped varinfo.
+└ @ DynamicPPLJETExt ~/Projects/public/DynamicPPL.jl/ext/DynamicPPLJETExt.jl:49
+
+julia> varinfo isa typeof(DynamicPPL.untyped_varinfo(model))
+true
+
+julia> # In contrast, a simple model with no random support can be handled by typed varinfo.
+       @model model_with_static_support() = x ~ Normal()
+model_with_static_support (generic function with 2 method)
+
+julia> varinfo = determine_suitable_varinfo(model_with_static_support());
+
+julia> varinfo isa typeof(DynamicPPL.typed_varinfo(model_with_static_support()))
+true
+```
+"""
+function determine_suitable_varinfo(
+    model::DynamicPPL.Model,
+    context::DynamicPPL.AbstractContext=DynamicPPL.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).
+        DynamicPPL.typed_varinfo(model, context)
+    end
+end
+
+end

src/sampler.jl

@@ -67,6 +67,20 @@ function AbstractMCMC.step(
     return vi, nothing
 end
 
+"""
+    default_varinfo(rng, model, sampler[, context])
+
+Return a default varinfo object for the given `model` and `sampler`.
+
+# Arguments
+- `rng::Random.AbstractRNG`: Random number generator.
+- `model::Model`: Model for which we want to create a varinfo object.
+- `sampler::AbstractSampler`: Sampler which will make use of the varinfo object.
+- `context::AbstractContext`: Context in which the model is evaluated.
+
+# Returns
+- `AbstractVarInfo`: Default varinfo object for the given `model` and `sampler`.
+"""
 function default_varinfo(rng::Random.AbstractRNG, model::Model, sampler::AbstractSampler)
     return default_varinfo(rng, model, sampler, DefaultContext())
 end
@@ -126,7 +140,7 @@ By default, `data` is returned.
 loadstate(data) = data
 
 """
-    default_chaintype(sampler)
+    default_chain_type(sampler)
 
 Default type of the chain of posterior samples from `sampler`.
 """