Merge branch 'breaking' into mhauru/accumulators-stage2

Author: penelopeysm

HISTORY.md

@@ -4,6 +4,10 @@
 
 **Breaking changes**
 
+### Submodel macro
+
+The `@submodel` macro is fully removed; please use `to_submodel` instead.
+
 ### Accumulators
 
 This release overhauls how VarInfo objects track variables such as the log joint probability. The new approach is to use what we call accumulators: Objects that the VarInfo carries on it that may change their state at each `tilde_assume!!` and `tilde_observe!!` call based on the value of the variable in question. They replace both variables that were previously hard-coded in the `VarInfo` object (`logp` and `num_produce`) and some contexts. This brings with it a number of breaking changes:
@@ -27,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.
@@ -50,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

@@ -146,24 +146,12 @@ to_submodel
 
 Note that a `[to_submodel](@ref)` is only sampleable; one cannot compute `logpdf` for its realizations.
 
-In the past, one would instead embed sub-models using [`@submodel`](@ref), which has been deprecated since the introduction of [`to_submodel(model)`](@ref)
-
-```@docs
-@submodel
-```
-
 In the context of including models within models, it's also useful to prefix the variables in sub-models to avoid variable names clashing:
 
 ```@docs
 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

@@ -128,7 +128,6 @@ export AbstractVarInfo,
     to_submodel,
     # Convenience macros
     @addlogprob!,
-    @submodel,
     value_iterator_from_chain,
     check_model,
     check_model_and_trace,
@@ -176,6 +175,7 @@ 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")
@@ -186,7 +186,6 @@ include("simple_varinfo.jl")
 include("context_implementations.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")
-include("submodel_macro.jl")
 include("transforming.jl")
 include("logdensityfunction.jl")
 include("model_utils.jl")
@@ -227,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