WIP other changes

Author: penelopeysm

ext/DynamicPPLJETExt.jl

@@ -4,15 +4,10 @@ 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,
+    model::DynamicPPL.Model, 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
-    )
+    f, argtypes = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(model, varinfo)
     # 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
@@ -24,14 +19,19 @@ function DynamicPPL.Experimental.is_suitable_varinfo(
 end
 
 function DynamicPPL.Experimental._determine_varinfo_jet(
-    model::DynamicPPL.Model, context::DynamicPPL.AbstractContext; only_ddpl::Bool=true
+    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(model, context)
+    varinfo = DynamicPPL.typed_varinfo(sampling_model)
 
     # 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
+        model, varinfo; only_ddpl
     )
 
     if !issuccess
@@ -46,7 +46,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(model, context)
+        DynamicPPL.untyped_varinfo(model)
     end
 end
 

src/compiler.jl

@@ -1,4 +1,4 @@
-const INTERNALNAMES = (:__model__, :__context__, :__varinfo__)
+const INTERNALNAMES = (:__model__, :__varinfo__)
 
 """
     need_concretize(expr)
@@ -63,9 +63,9 @@ used in its place.
 function isassumption(expr::Union{Expr,Symbol}, vn=make_varname_expression(expr))
     return quote
         if $(DynamicPPL.contextual_isassumption)(
-            __context__, $(DynamicPPL.prefix)(__context__, $vn)
+            __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
         )
-            # Considered an assumption by `__context__` which means either:
+            # Considered an assumption by `__model__.context` which means either:
             # 1. We hit the default implementation, e.g. using `DefaultContext`,
             #    which in turn means that we haven't considered if it's one of
             #    the model arguments, hence we need to check this.
@@ -116,7 +116,7 @@ end
 isfixed(expr, vn) = false
 function isfixed(::Union{Symbol,Expr}, vn)
     return :($(DynamicPPL.contextual_isfixed)(
-        __context__, $(DynamicPPL.prefix)(__context__, $vn)
+        __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
     ))
 end
 
@@ -417,7 +417,7 @@ function generate_assign(left, right)
     return quote
         $right_val = $right
         if $(DynamicPPL.is_extracting_values)(__varinfo__)
-            $vn = $(DynamicPPL.prefix)(__context__, $(make_varname_expression(left)))
+            $vn = $(DynamicPPL.prefix)(__model__.context, $(make_varname_expression(left)))
             __varinfo__ = $(map_accumulator!!)(
                 $acc -> push!($acc, $vn, $right_val), __varinfo__, Val(:ValuesAsInModel)
             )
@@ -431,7 +431,11 @@ function generate_tilde_literal(left, right)
     @gensym value
     return quote
         $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
-            __context__, $(DynamicPPL.check_tilde_rhs)($right), $left, nothing, __varinfo__
+            __model__.context,
+            $(DynamicPPL.check_tilde_rhs)($right),
+            $left,
+            nothing,
+            __varinfo__,
         )
         $value
     end
@@ -456,20 +460,20 @@ function generate_tilde(left, right)
         $isassumption = $(DynamicPPL.isassumption(left, vn))
         if $(DynamicPPL.isfixed(left, vn))
             $left = $(DynamicPPL.getfixed_nested)(
-                __context__, $(DynamicPPL.prefix)(__context__, $vn)
+                __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
             )
         elseif $isassumption
             $(generate_tilde_assume(left, dist, vn))
         else
             # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
             if !$(DynamicPPL.inargnames)($vn, __model__)
                 $left = $(DynamicPPL.getconditioned_nested)(
-                    __context__, $(DynamicPPL.prefix)(__context__, $vn)
+                    __model__.context, $(DynamicPPL.prefix)(__model__.context, $vn)
                 )
             end
 
             $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(
-                __context__,
+                __model__.context,
                 $(DynamicPPL.check_tilde_rhs)($dist),
                 $(maybe_view(left)),
                 $vn,
@@ -494,7 +498,7 @@ function generate_tilde_assume(left, right, vn)
 
     return quote
         $value, __varinfo__ = $(DynamicPPL.tilde_assume!!)(
-            __context__,
+            __model__.context,
             $(DynamicPPL.unwrap_right_vn)($(DynamicPPL.check_tilde_rhs)($right), $vn)...,
             __varinfo__,
         )
@@ -652,11 +656,7 @@ function build_output(modeldef, linenumbernode)
 
     # Add the internal arguments to the user-specified arguments (positional + keywords).
     evaluatordef[:args] = vcat(
-        [
-            :(__model__::$(DynamicPPL.Model)),
-            :(__varinfo__::$(DynamicPPL.AbstractVarInfo)),
-            :(__context__::$(DynamicPPL.AbstractContext)),
-        ],
+        [:(__model__::$(DynamicPPL.Model)), :(__varinfo__::$(DynamicPPL.AbstractVarInfo))],
         args,
     )
 

src/debug_utils.jl

@@ -535,14 +535,13 @@ function has_static_constraints(
 end
 
 """
-    gen_evaluator_call_with_types(model[, varinfo, context])
+    gen_evaluator_call_with_types(model[, varinfo])
 
 Generate the evaluator call and the types of the arguments.
 
 # Arguments
 - `model::Model`: The model whose evaluator is of interest.
 - `varinfo::AbstractVarInfo`: The varinfo to use when evaluating the model. Default: `VarInfo(model)`.
-- `context::AbstractContext`: The context to use when evaluating the model. Default: [`DefaultContext`](@ref).
 
 # Returns
 A 2-tuple with the following elements:
@@ -551,11 +550,9 @@ A 2-tuple with the following elements:
 - `argtypes::Type{<:Tuple}`: The types of the arguments for the evaluator.
 """
 function gen_evaluator_call_with_types(
-    model::Model,
-    varinfo::AbstractVarInfo=VarInfo(model),
-    context::AbstractContext=DefaultContext(),
+    model::Model, varinfo::AbstractVarInfo=VarInfo(model)
 )
-    args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
+    args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo)
     return if isempty(kwargs)
         (model.f, Base.typesof(args...))
     else
@@ -564,7 +561,7 @@ function gen_evaluator_call_with_types(
 end
 
 """
-    model_warntype(model[, varinfo, context]; optimize=true)
+    model_warntype(model[, varinfo]; optimize=true)
 
 Check the type stability of the model's evaluator, warning about any potential issues.
 
@@ -573,23 +570,19 @@ This simply calls `@code_warntype` on the model's evaluator, filling in internal
 # Arguments
 - `model::Model`: The model to check.
 - `varinfo::AbstractVarInfo`: The varinfo to use when evaluating the model. Default: `VarInfo(model)`.
-- `context::AbstractContext`: The context to use when evaluating the model. Default: [`DefaultContext`](@ref).
 
 # Keyword Arguments
 - `optimize::Bool`: Whether to generate optimized code. Default: `false`.
 """
 function model_warntype(
-    model::Model,
-    varinfo::AbstractVarInfo=VarInfo(model),
-    context::AbstractContext=DefaultContext();
-    optimize::Bool=false,
+    model::Model, varinfo::AbstractVarInfo=VarInfo(model), optimize::Bool=false
 )
-    ftype, argtypes = gen_evaluator_call_with_types(model, varinfo, context)
+    ftype, argtypes = gen_evaluator_call_with_types(model, varinfo)
     return InteractiveUtils.code_warntype(ftype, argtypes; optimize=optimize)
 end
 
 """
-    model_typed(model[, varinfo, context]; optimize=true)
+    model_typed(model[, varinfo]; optimize=true)
 
 Return the type inference for the model's evaluator.
 
@@ -598,18 +591,14 @@ This simply calls `@code_typed` on the model's evaluator, filling in internal ar
 # Arguments
 - `model::Model`: The model to check.
 - `varinfo::AbstractVarInfo`: The varinfo to use when evaluating the model. Default: `VarInfo(model)`.
-- `context::AbstractContext`: The context to use when evaluating the model. Default: [`DefaultContext`](@ref).
 
 # Keyword Arguments
 - `optimize::Bool`: Whether to generate optimized code. Default: `true`.
 """
 function model_typed(
-    model::Model,
-    varinfo::AbstractVarInfo=VarInfo(model),
-    context::AbstractContext=DefaultContext();
-    optimize::Bool=true,
+    model::Model, varinfo::AbstractVarInfo=VarInfo(model), optimize::Bool=true
 )
-    ftype, argtypes = gen_evaluator_call_with_types(model, varinfo, context)
+    ftype, argtypes = gen_evaluator_call_with_types(model, varinfo)
     return only(InteractiveUtils.code_typed(ftype, argtypes; optimize=optimize))
 end
 

src/experimental.jl

@@ -4,16 +4,15 @@ using DynamicPPL: DynamicPPL
 
 # This file only defines the names of the functions, and their docstrings. The actual implementations are in `ext/DynamicPPLJETExt.jl`, since we don't want to depend on JET.jl other than as a weak dependency.
 """
-    is_suitable_varinfo(model::Model, context::AbstractContext, varinfo::AbstractVarInfo; kwargs...)
+    is_suitable_varinfo(model::Model, varinfo::AbstractVarInfo; kwargs...)
 
-Check if the `model` supports evaluation using the provided `context` and `varinfo`.
+Check if the `model` supports evaluation using the provided `varinfo`.
 
 !!! warning
     Loading JET.jl is required before calling this function.
 
 # Arguments
 - `model`: The model to verify the support for.
-- `context`: The context to use for the model evaluation.
 - `varinfo`: The varinfo to verify the support for.
 
 # Keyword Arguments
@@ -29,7 +28,7 @@ function is_suitable_varinfo end
 function _determine_varinfo_jet end
 
 """
-    determine_suitable_varinfo(model[, context]; only_ddpl::Bool=true)
+    determine_suitable_varinfo(model; only_ddpl::Bool=true)
 
 Return a suitable varinfo for the given `model`.
 
@@ -41,7 +40,6 @@ See also: [`DynamicPPL.Experimental.is_suitable_varinfo`](@ref).
 
 # 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.
@@ -85,14 +83,10 @@ julia> vi 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,
-)
+function determine_suitable_varinfo(model::DynamicPPL.Model, 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)
+        _determine_varinfo_jet(model; only_ddpl)
     else
         # Warn the user.
         @warn "JET.jl is not loaded. Assumes the model is compatible with typed varinfo."

src/extract_priors.jl

@@ -116,7 +116,8 @@ function extract_priors(rng::Random.AbstractRNG, model::Model)
     # workaround for the fact that `order` is still hardcoded in VarInfo, and hence you
     # can't push new variables without knowing the num_produce. Remove this when possible.
     varinfo = setaccs!!(varinfo, (PriorDistributionAccumulator(), NumProduceAccumulator()))
-    varinfo = last(evaluate!!(model, varinfo, SamplingContext(rng)))
+    new_model = contextualize(model, SamplingContext(rng, model.context))
+    varinfo = last(evaluate!!(new_model, varinfo))
     return getacc(varinfo, Val(:PriorDistributionAccumulator)).priors
 end
 
@@ -135,6 +136,6 @@ function extract_priors(model::Model, varinfo::AbstractVarInfo)
     varinfo = setaccs!!(
         deepcopy(varinfo), (PriorDistributionAccumulator(), NumProduceAccumulator())
     )
-    varinfo = last(evaluate!!(model, varinfo, DefaultContext()))
+    varinfo = last(evaluate!!(model, varinfo))
     return getacc(varinfo, Val(:PriorDistributionAccumulator)).priors
 end