Fix literally everything else that I broke

Author: penelopeysm

benchmarks/src/DynamicPPLBenchmarks.jl

@@ -81,13 +81,12 @@ function make_suite(model, varinfo_choice::Symbol, adbackend::Symbol, islinked::
     end
 
     adbackend = to_backend(adbackend)
-    context = DynamicPPL.DefaultContext()
 
     if islinked
         vi = DynamicPPL.link(vi, model)
     end
 
-    f = DynamicPPL.LogDensityFunction(model, vi, context; adtype=adbackend)
+    f = DynamicPPL.LogDensityFunction(model, vi; adtype=adbackend)
     # The parameters at which we evaluate f.
     θ = vi[:]
 

docs/src/api.md

@@ -36,6 +36,12 @@ getargnames
 getmissings
 ```
 
+The context of a model can be set using [`contextualize`](@ref):
+
+```@docs
+contextualize
+```
+
 ## Evaluation
 
 With [`rand`](@ref) one can draw samples from the prior distribution of a [`Model`](@ref).
@@ -438,13 +444,21 @@ DynamicPPL.varname_and_value_leaves
 
 ### Evaluation Contexts
 
-Internally, both sampling and evaluation of log densities are performed with [`AbstractPPL.evaluate!!`](@ref).
+Internally, model evaluation is performed with [`AbstractPPL.evaluate!!`](@ref).
 
 ```@docs
 AbstractPPL.evaluate!!
 ```
 
-The behaviour of a model execution can be changed with evaluation contexts that are passed as additional argument to the model function.
+This method mutates the `varinfo` used for execution.
+By default, it does not perform any actual sampling: it only evaluates the model using the values of the variables that are already in the `varinfo`.
+To perform sampling, you can either wrap `model.context` in a `SamplingContext`, or use this convenience method:
+
+```@docs
+DynamicPPL.sample!!
+```
+
+The behaviour of a model execution can be changed with evaluation contexts, which are a field of the model.
 Contexts are subtypes of `AbstractPPL.AbstractContext`.
 
 ```@docs

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
+        sampling_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(sampling_model)
     end
 end
 

ext/DynamicPPLMCMCChainsExt.jl

@@ -115,7 +115,7 @@ function DynamicPPL.predict(
     iters = Iterators.product(1:size(chain, 1), 1:size(chain, 3))
     predictive_samples = map(iters) do (sample_idx, chain_idx)
         DynamicPPL.setval_and_resample!(varinfo, parameter_only_chain, sample_idx, chain_idx)
-        model(rng, varinfo, DynamicPPL.SampleFromPrior())
+        varinfo = last(DynamicPPL.sample!!(rng, model, varinfo))
 
         vals = DynamicPPL.values_as_in_model(model, false, varinfo)
         varname_vals = mapreduce(

src/DynamicPPL.jl

@@ -102,6 +102,7 @@ export AbstractVarInfo,
     # LogDensityFunction
     LogDensityFunction,
     # Contexts
+    contextualize,
     SamplingContext,
     DefaultContext,
     PrefixContext,

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

@@ -131,9 +131,7 @@ A context used for checking validity of a model.
 # Fields
 $(FIELDS)
 """
-struct DebugContext{M<:Model,C<:AbstractContext} <: AbstractContext
-    "model that is being run"
-    model::M
+struct DebugContext{C<:AbstractContext} <: AbstractContext
     "context used for running the model"
     context::C
     "mapping from varnames to the number of times they have been seen"
@@ -149,7 +147,6 @@ struct DebugContext{M<:Model,C<:AbstractContext} <: AbstractContext
 end
 
 function DebugContext(
-    model::Model,
     context::AbstractContext=DefaultContext();
     varnames_seen=OrderedDict{VarName,Int}(),
     statements=Vector{Stmt}(),
@@ -158,7 +155,6 @@ function DebugContext(
     record_varinfo=false,
 )
     return DebugContext(
-        model,
         context,
         varnames_seen,
         statements,
@@ -344,7 +340,7 @@ function check_varnames_seen(varnames_seen::AbstractDict{VarName,Int})
 end
 
 # A check we run on the model before evaluating it.
-function check_model_pre_evaluation(context::DebugContext, model::Model)
+function check_model_pre_evaluation(model::Model)
     issuccess = true
     # If something is in the model arguments, then it should NOT be in `condition`,
     # nor should there be any symbol present in `condition` that has the same symbol.
@@ -361,8 +357,8 @@ function check_model_pre_evaluation(context::DebugContext, model::Model)
     return issuccess
 end
 
-function check_model_post_evaluation(context::DebugContext, model::Model)
-    return check_varnames_seen(context.varnames_seen)
+function check_model_post_evaluation(model::Model)
+    return check_varnames_seen(model.context.varnames_seen)
 end
 
 """
@@ -438,25 +434,23 @@ function check_model_and_trace(
     rng::Random.AbstractRNG,
     model::Model;
     varinfo=VarInfo(),
-    context=SamplingContext(rng),
     error_on_failure=false,
     kwargs...,
 )
     # Execute the model with the debug context.
     debug_context = DebugContext(
-        model, context; error_on_failure=error_on_failure, kwargs...
+        SamplingContext(rng, model.context); error_on_failure=error_on_failure, kwargs...
     )
+    debug_model = DynamicPPL.contextualize(model, debug_context)
 
     # Perform checks before evaluating the model.
-    issuccess = check_model_pre_evaluation(debug_context, model)
+    issuccess = check_model_pre_evaluation(debug_model)
 
     # Force single-threaded execution.
-    retval, varinfo_result = DynamicPPL.evaluate_threadunsafe!!(
-        model, varinfo, debug_context
-    )
+    DynamicPPL.evaluate_threadunsafe!!(debug_model, varinfo)
 
     # Perform checks after evaluating the model.
-    issuccess &= check_model_post_evaluation(debug_context, model)
+    issuccess &= check_model_post_evaluation(debug_model)
 
     if !issuccess && error_on_failure
         error("model check failed")
@@ -535,14 +529,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 +544,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 +555,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 +564,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 +585,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