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

docs/src/internals/submodel_condition.md

@@ -102,13 +102,13 @@ Take these models, for example:
 unwrap_sampling_context(ctx::DynamicPPL.SamplingContext) = ctx.context
 unwrap_sampling_context(ctx::DynamicPPL.AbstractContext) = ctx
 @model function inner()
-    println("inner context: $(unwrap_sampling_context(__context__))")
+    println("inner context: $(unwrap_sampling_context(__model__.context))")
     x ~ Normal()
     return y ~ Normal()
 end
 
 @model function outer()
-    println("outer context: $(unwrap_sampling_context(__context__))")
+    println("outer context: $(unwrap_sampling_context(__model__.context))")
     return a ~ to_submodel(inner())
 end
 
@@ -118,7 +118,7 @@ with_outer_cond = outer() | (@varname(a.x) => 1.0)
 # 'Inner conditioning'
 inner_cond = inner() | (@varname(x) => 1.0)
 @model function outer2()
-    println("outer context: $(unwrap_sampling_context(__context__))")
+    println("outer context: $(unwrap_sampling_context(__model__.context))")
     return a ~ to_submodel(inner_cond)
 end
 with_inner_cond = outer2()

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,
     )