Tidy up tilde-pipeline methods and docstrings

Author: penelopeysm

src/context_implementations.jl

@@ -1,4 +1,13 @@
-# assume
+"""
+    tilde_assume!!(context, right::Distribution, vn, vi)
+
+Handle assumed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
+accumulate the log probability, and return the sampled value and updated `vi`.
+
+    tilde_assume!!(context, right::DynamicPPL.Submodel, vn, vi)
+
+Evaluate the submodel with the given context.
+"""
 function tilde_assume!!(context::AbstractContext, right::Distribution, vn, vi)
     return tilde_assume!!(childcontext(context), right, vn, vi)
 end
@@ -22,24 +31,23 @@ function tilde_assume!!(context::PrefixContext, right::Distribution, vn, vi)
     new_vn, new_context = prefix_and_strip_contexts(context, vn)
     return tilde_assume!!(new_context, right, new_vn, vi)
 end
+function tilde_assume!!(context::AbstractContext, right::DynamicPPL.Submodel, vn, vi)
+    return _evaluate!!(right, vi, context, vn)
+end
 
 """
-    tilde_assume!!(context, right, vn, vi)
+    tilde_observe!!(context, right::Distribution, left, vn, vi)
 
-Handle assumed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
-accumulate the log probability, and return the sampled value and updated `vi`.
-"""
-function tilde_assume!!(context, right::DynamicPPL.Submodel, vn, vi)
-    return _evaluate!!(right, vi, context, vn)
-end
+Handle observed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
+accumulate the log probability, and return the observed value and updated `vi`.
 
-# observe
-function tilde_observe!!(context::AbstractContext, right, left, vn, 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::AbstractContext, right::Distribution, left, vn, vi)
     return tilde_observe!!(childcontext(context), right, left, vn, vi)
 end
-
-# `PrefixContext`
-function tilde_observe!!(context::PrefixContext, right, left, vn, vi)
+function tilde_observe!!(context::PrefixContext, right::Distribution, left, vn, vi)
     # In the observe case, unlike assume, `vn` may be `nothing` if the LHS is a literal
     # value. For the need for prefix_and_strip_contexts rather than just prefix, see the
     # comment in `tilde_assume!!`.
@@ -50,21 +58,10 @@ function tilde_observe!!(context::PrefixContext, right, left, vn, vi)
     end
     return tilde_observe!!(new_context, right, left, new_vn, vi)
 end
-
-"""
-    tilde_observe!!(context, right, left, vn, vi)
-
-Handle observed variables, e.g., `x ~ Normal()` (where `x` does occur in the model inputs),
-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!!(::DefaultContext, right::Distribution, left, vn, vi)
     vi = accumulate_observe!!(vi, right, left, vn)
     return left, vi
 end
-
-function tilde_observe!!(::DefaultContext, ::DynamicPPL.Submodel, left, vn, vi)
+function tilde_observe!!(::AbstractContext, ::DynamicPPL.Submodel, left, vn, vi)
     throw(ArgumentError("`x ~ to_submodel(...)` is not supported when `x` is observed"))
 end

src/contexts/init.jl

@@ -191,6 +191,12 @@ function tilde_assume!!(
     return x, vi
 end
 
-function tilde_observe!!(::InitContext, right, left, vn, vi)
+function tilde_observe!!(
+    ::InitContext,
+    right::Distribution,
+    left,
+    vn::Union{VarName,Nothing},
+    vi::AbstractVarInfo,
+)
     return tilde_observe!!(DefaultContext(), right, left, vn, vi)
 end