Remove 3-argument {_,}evaluate!!; clean up submodel code (#960)

* Clean up submodel code, remove 3-arg `_evaluate!!`

* Remove 3-argument `evaluate!!` as well

* Update changelog

* Improve submodel error message

* Fix doctest

* Add error hint for three-argument evaluate!!

Author: penelopeysm

HISTORY.md

@@ -31,11 +31,11 @@ This version therefore excises the context argument, and instead uses `model.con
 The upshot of this is that many functions that previously took a context argument now no longer do.
 There were very few such functions where the context argument was actually used (most of them simply took `DefaultContext()` as the default value).
 
-`evaluate!!(model, varinfo, ext_context)` is deprecated, and broadly speaking you should replace calls to that with `new_model = contextualize(model, ext_context); evaluate!!(new_model, varinfo)`.
+`evaluate!!(model, varinfo, ext_context)` is removed, and broadly speaking you should replace calls to that with `new_model = contextualize(model, ext_context); evaluate!!(new_model, varinfo)`.
 If the 'external context' `ext_context` is a parent context, then you should wrap `model.context` appropriately to ensure that its information content is not lost.
 If, on the other hand, `ext_context` is a `DefaultContext`, then you can just drop the argument entirely.
 
-To aid with this process, `contextualize` is now exported from DynamicPPL.
+**To aid with this process, `contextualize` is now exported from DynamicPPL.**
 
 The main situation where one _did_ want to specify an additional evaluation context was when that context was a `SamplingContext`.
 Doing this would allow you to run the model and sample fresh values, instead of just using the values that existed in the VarInfo object.
@@ -54,9 +54,10 @@ However, here are the more user-facing ones:
 
 And a couple of more internal changes:
 
-  - `evaluate!!`, `evaluate_threadsafe!!`, and `evaluate_threadunsafe!!` no longer accept context arguments
+  - Just like `evaluate!!`, the other functions `_evaluate!!`, `evaluate_threadsafe!!`, and `evaluate_threadunsafe!!` now no longer accept context arguments
   - `evaluate!!` no longer takes rng and sampler (if you used this, you should use `evaluate_and_sample!!` instead, or construct your own `SamplingContext`)
   - The model evaluation function, `model.f` for some `model::Model`, no longer takes a context as an argument
+  - The internal representation and API dealing with submodels (i.e., `ReturnedModelWrapper`, `Sampleable`, `should_auto_prefix`, `is_rhs_model`) has been simplified. If you need to check whether something is a submodel, just use `x isa DynamicPPL.Submodel`. Note that the public API i.e. `to_submodel` remains completely untouched.
 
 ## 0.36.12
 

docs/src/api.md

@@ -152,12 +152,6 @@ In the context of including models within models, it's also useful to prefix the
 DynamicPPL.prefix
 ```
 
-Under the hood, [`to_submodel`](@ref) makes use of the following method to indicate that the model it's wrapping is a model over its return-values rather than something else
-
-```@docs
-returned(::Model)
-```
-
 ## Utilities
 
 It is possible to manually increase (or decrease) the accumulated log likelihood or prior from within a model function.

src/DynamicPPL.jl

@@ -171,11 +171,11 @@ abstract type AbstractVarInfo <: AbstractModelTrace end
 include("utils.jl")
 include("chains.jl")
 include("model.jl")
-include("submodel.jl")
 include("sampler.jl")
 include("varname.jl")
 include("distribution_wrappers.jl")
 include("contexts.jl")
+include("submodel.jl")
 include("varnamedvector.jl")
 include("accumulators.jl")
 include("default_accumulators.jl")
@@ -226,6 +226,21 @@ if isdefined(Base.Experimental, :register_error_hint)
                 )
             end
         end
+
+        Base.Experimental.register_error_hint(MethodError) do io, exc, argtypes, _
+            is_evaluate_three_arg =
+                exc.f === AbstractPPL.evaluate!! &&
+                length(argtypes) == 3 &&
+                argtypes[1] <: Model &&
+                argtypes[2] <: AbstractVarInfo &&
+                argtypes[3] <: AbstractContext
+            if is_evaluate_three_arg
+                print(
+                    io,
+                    "\n\nThe method `evaluate!!(model, varinfo, new_ctx)` has been removed. Instead, you should store the `new_ctx` in the `model.context` field using `new_model = contextualize(model, new_ctx)`, and then call `evaluate!!(new_model, varinfo)` on the new model. (Note that, if the model already contained a non-default context, you will need to wrap the existing context.)",
+                )
+            end
+        end
     end
 end
 

src/compiler.jl

@@ -176,11 +176,7 @@ function check_tilde_rhs(@nospecialize(x))
 end
 check_tilde_rhs(x::Distribution) = x
 check_tilde_rhs(x::AbstractArray{<:Distribution}) = x
-check_tilde_rhs(x::ReturnedModelWrapper) = x
-function check_tilde_rhs(x::Sampleable{<:Any,AutoPrefix}) where {AutoPrefix}
-    model = check_tilde_rhs(x.model)
-    return Sampleable{typeof(model),AutoPrefix}(model)
-end
+check_tilde_rhs(x::Submodel{M,AutoPrefix}) where {M,AutoPrefix} = x
 
 """
     check_dot_tilde_rhs(x)

src/context_implementations.jl

@@ -63,31 +63,10 @@ By default, calls `tilde_assume(context, right, vn, vi)` and accumulates the log
 probability of `vi` with the returned value.
 """
 function tilde_assume!!(context, right, vn, vi)
-    return if is_rhs_model(right)
-        # Here, we apply the PrefixContext _not_ to the parent `context`, but
-        # to the context of the submodel being evaluated. This means that later=
-        # on in `make_evaluate_args_and_kwargs`, the context stack will be
-        # correctly arranged such that it goes like this:
-        #  parent_context[1] -> parent_context[2] -> ... -> PrefixContext ->
-        #    submodel_context[1] -> submodel_context[2] -> ... -> leafcontext
-        # See the docstring of `make_evaluate_args_and_kwargs`, and the internal
-        # DynamicPPL documentation on submodel conditioning, for more details.
-        #
-        # NOTE: This relies on the existence of `right.model.model`. Right now,
-        # the only thing that can return true for `is_rhs_model` is something
-        # (a `Sampleable`) that has a `model` field that itself (a
-        # `ReturnedModelWrapper`) has a `model` field. This may or may not
-        # change in the future.
-        if should_auto_prefix(right)
-            dppl_model = right.model.model # This isa DynamicPPL.Model
-            prefixed_submodel_context = PrefixContext(vn, dppl_model.context)
-            new_dppl_model = contextualize(dppl_model, prefixed_submodel_context)
-            right = to_submodel(new_dppl_model, true)
-        end
-        rand_like!!(right, context, vi)
+    return if right isa DynamicPPL.Submodel
+        _evaluate!!(right, vi, context, vn)
     else
-        value, vi = tilde_assume(context, right, vn, vi)
-        return value, vi
+        tilde_assume(context, right, vn, vi)
     end
 end
 
@@ -129,17 +108,14 @@ accumulate the log probability, and return the observed value and updated `vi`.
 Falls back to `tilde_observe!!(context, right, left, vi)` ignoring the information about variable name
 and indices; if needed, these can be accessed through this function, though.
 """
-function tilde_observe!!(context::DefaultContext, right, left, vn, vi)
-    is_rhs_model(right) && throw(
-        ArgumentError(
-            "`~` with a model on the right-hand side of an observe statement is not supported",
-        ),
-    )
+function tilde_observe!!(::DefaultContext, right, left, vn, vi)
+    right isa DynamicPPL.Submodel &&
+        throw(ArgumentError("`x ~ to_submodel(...)` is not supported when `x` is observed"))
     vi = accumulate_observe!!(vi, right, left, vn)
     return left, vi
 end
 
-function assume(rng::Random.AbstractRNG, spl::Sampler, dist)
+function assume(::Random.AbstractRNG, spl::Sampler, dist)
     return error("DynamicPPL.assume: unmanaged inference algorithm: $(typeof(spl))")
 end
 

src/model.jl

@@ -258,7 +258,7 @@ julia> # However, it's not possible to condition `inner` directly.
        conditioned_model_fail = model | (inner = 1.0, );
 
 julia> conditioned_model_fail()
-ERROR: ArgumentError: `~` with a model on the right-hand side of an observe statement is not supported
+ERROR: ArgumentError: `x ~ to_submodel(...)` is not supported when `x` is observed
 [...]
 ```
 """
@@ -864,12 +864,6 @@ If multiple threads are available, the varinfo provided will be wrapped in a
 
 Returns a tuple of the model's return value, plus the updated `varinfo`
 (unwrapped if necessary).
-
-    evaluate!!(model::Model, varinfo, context)
-
-When an extra context stack is provided, the model's context is inserted into
-that context stack. See `combine_model_and_external_contexts`. This method is
-deprecated.
 """
 function AbstractPPL.evaluate!!(model::Model, varinfo::AbstractVarInfo)
     return if use_threadsafe_eval(model.context, varinfo)
@@ -878,17 +872,6 @@ function AbstractPPL.evaluate!!(model::Model, varinfo::AbstractVarInfo)
         evaluate_threadunsafe!!(model, varinfo)
     end
 end
-function AbstractPPL.evaluate!!(
-    model::Model, varinfo::AbstractVarInfo, context::AbstractContext
-)
-    Base.depwarn(
-        "The `context` argument to evaluate!!(model, varinfo, context) is deprecated.",
-        :dynamicppl_evaluate_context,
-    )
-    new_ctx = combine_model_and_external_contexts(model.context, context)
-    model = contextualize(model, new_ctx)
-    return evaluate!!(model, varinfo)
-end
 
 """
     evaluate_threadunsafe!!(model, varinfo)
@@ -932,54 +915,14 @@ Evaluate the `model` with the given `varinfo`.
 
 This function does not wrap the varinfo in a `ThreadSafeVarInfo`. It also does not
 reset the log probability of the `varinfo` before running.
-
-    _evaluate!!(model::Model, varinfo, context)
-
-If an additional `context` is provided, the model's context is combined with
-that context before evaluation.
 """
 function _evaluate!!(model::Model, varinfo::AbstractVarInfo)
     args, kwargs = make_evaluate_args_and_kwargs(model, varinfo)
     return model.f(args...; kwargs...)
 end
-function _evaluate!!(model::Model, varinfo::AbstractVarInfo, context::AbstractContext)
-    # TODO(penelopeysm): We don't really need this, but it's a useful
-    # convenience method. We could remove it after we get rid of the
-    # evaluate_threadsafe!! stuff (in favour of making users call evaluate!!
-    # with a TSVI themselves).
-    new_ctx = combine_model_and_external_contexts(model.context, context)
-    model = contextualize(model, new_ctx)
-    return _evaluate!!(model, varinfo)
-end
 
 is_splat_symbol(s::Symbol) = startswith(string(s), "#splat#")
 
-"""
-    combine_model_and_external_contexts(model_context, external_context)
-
-Combine a context from a model and an external context into a single context.
-
-The resulting context stack has the following structure:
-
-    `external_context` -> `childcontext(external_context)` -> ... ->
-    `model_context` -> `childcontext(model_context)` -> ... ->
-    `leafcontext(external_context)`
-
-The reason for this is that we want to give `external_context` precedence over
-`model_context`, while also preserving the leaf context of `external_context`.
-We can do this by
-
-1. Set the leaf context of `model_context` to `leafcontext(external_context)`.
-2. Set leaf context of `external_context` to the context resulting from (1).
-"""
-function combine_model_and_external_contexts(
-    model_context::AbstractContext, external_context::AbstractContext
-)
-    return setleafcontext(
-        external_context, setleafcontext(model_context, leafcontext(external_context))
-    )
-end
-
 """
     make_evaluate_args_and_kwargs(model, varinfo)