Add setleafcontext(::Model, ::AbstractContext)

Author: penelopeysm

HISTORY.md

@@ -52,6 +52,10 @@ The separation of these functions was primarily implemented to avoid performing
 
 **Other changes**
 
+### `setleafcontext(model, context)`
+
+This convenience method has been added to quickly modify the leaf context of a model.
+
 ### Reimplementation of functions using `InitContext`
 
 A number of functions have been reimplemented and unified with the help of `InitContext`.

ext/DynamicPPLJETExt.jl

@@ -24,9 +24,7 @@ function DynamicPPL.Experimental._determine_varinfo_jet(
     varinfo = DynamicPPL.typed_varinfo(model)
 
     # Check type stability of evaluation (i.e. DefaultContext)
-    model = DynamicPPL.contextualize(
-        model, DynamicPPL.setleafcontext(model.context, DynamicPPL.DefaultContext())
-    )
+    model = DynamicPPL.setleafcontext(model, DynamicPPL.DefaultContext())
     eval_issuccess, eval_result = DynamicPPL.Experimental.is_suitable_varinfo(
         model, varinfo; only_ddpl
     )
@@ -36,9 +34,7 @@ function DynamicPPL.Experimental._determine_varinfo_jet(
     end
 
     # Check type stability of initialisation (i.e. InitContext)
-    model = DynamicPPL.contextualize(
-        model, DynamicPPL.setleafcontext(model.context, DynamicPPL.InitContext())
-    )
+    model = DynamicPPL.setleafcontext(model, DynamicPPL.InitContext())
     init_issuccess, init_result = DynamicPPL.Experimental.is_suitable_varinfo(
         model, varinfo; only_ddpl
     )

src/model.jl

@@ -94,6 +94,15 @@ with its underlying context set to `context`.
 function contextualize(model::Model, context::AbstractContext)
     return Model(model.f, model.args, model.defaults, context)
 end
+"""
+    setleafcontext(model::Model, context::AbstractContext)
+
+Return a new `Model` with its leaf context set to `context`. This is a convenience
+shortcut for `contextualize(model, setleafcontext(model.context, context)`).
+"""
+function setleafcontext(model::Model, context::AbstractContext)
+    return contextualize(model, setleafcontext(model.context, context))
+end
 
 """
     model | (x = 1.0, ...)
@@ -886,8 +895,7 @@ function init!!(
     varinfo::AbstractVarInfo,
     init_strategy::AbstractInitStrategy=InitFromPrior(),
 )
-    new_context = setleafcontext(model.context, InitContext(rng, init_strategy))
-    new_model = contextualize(model, new_context)
+    new_model = setleafcontext(model, InitContext(rng, init_strategy))
     return evaluate!!(new_model, varinfo)
 end
 function init!!(

src/test_utils/contexts.jl

@@ -47,7 +47,7 @@ function test_leaf_context(context::DynamicPPL.AbstractContext, model::DynamicPP
         _, untyped_vi = DynamicPPL.init!!(model, DynamicPPL.VarInfo())
         typed_vi = DynamicPPL.typed_varinfo(untyped_vi)
         # Set the test context as the new leaf context
-        new_model = contextualize(model, DynamicPPL.setleafcontext(model.context, context))
+        new_model = DynamicPPL.setleafcontext(model, context)
         # Check that evaluation works
         for vi in [untyped_vi, typed_vi]
             _, vi = DynamicPPL.evaluate!!(new_model, vi)

src/threadsafe.jl

@@ -103,16 +103,12 @@ end
 # consistency between `vi.accs_by_thread` field and `getacc(vi.varinfo)`, which accumulates
 # to define `getacc(vi)`.
 function link!!(t::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
-    model = contextualize(
-        model, setleafcontext(model.context, DynamicTransformationContext{false}())
-    )
+    model = setleafcontext(model, DynamicTransformationContext{false}())
     return settrans!!(last(evaluate!!(model, vi)), t)
 end
 
 function invlink!!(::DynamicTransformation, vi::ThreadSafeVarInfo, model::Model)
-    model = contextualize(
-        model, setleafcontext(model.context, DynamicTransformationContext{true}())
-    )
+    model = setleafcontext(model, DynamicTransformationContext{true}())
     return settrans!!(last(evaluate!!(model, vi)), NoTransformation())
 end
 

src/transforming.jl

@@ -59,7 +59,7 @@ function _transform!!(
     model::Model,
 )
     # To transform using DynamicTransformationContext, we evaluate the model using that as the leaf context:
-    model = contextualize(model, setleafcontext(model.context, ctx))
+    model = setleafcontext(model, ctx)
     vi = settrans!!(last(evaluate!!(model, vi)), t)
     return vi
 end

test/ext/DynamicPPLJETExt.jl

@@ -81,20 +81,14 @@
                 typed_vi = DynamicPPL.typed_varinfo(model)
 
                 @info "Evaluating with DefaultContext:"
-                model = DynamicPPL.contextualize(
-                    model,
-                    DynamicPPL.setleafcontext(model.context, DynamicPPL.DefaultContext()),
-                )
+                model = DynamicPPL.setleafcontext(model, DynamicPPL.DefaultContext())
                 f, argtypes = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
                     model, varinfo
                 )
                 JET.test_call(f, argtypes)
 
                 @info "Initialising with InitContext:"
-                model = DynamicPPL.contextualize(
-                    model,
-                    DynamicPPL.setleafcontext(model.context, DynamicPPL.InitContext()),
-                )
+                model = DynamicPPL.setleafcontext(model, DynamicPPL.InitContext())
                 f, argtypes = DynamicPPL.DebugUtils.gen_evaluator_call_with_types(
                     model, varinfo
                 )