Remove tilde_assume as well

Author: penelopeysm

docs/src/api.md

@@ -516,9 +516,3 @@ There is also the _experimental_ [`DynamicPPL.Experimental.determine_suitable_va
 DynamicPPL.Experimental.determine_suitable_varinfo
 DynamicPPL.Experimental.is_suitable_varinfo
 ```
-
-### [Model-Internal Functions](@id model_internal)
-
-```@docs
-tilde_assume
-```

src/context_implementations.jl

@@ -1,43 +1,36 @@
 # assume
-function tilde_assume(context::AbstractContext, args...)
-    return tilde_assume(childcontext(context), args...)
+function tilde_assume!!(context::AbstractContext, right::Distribution, vn, vi)
+    return tilde_assume!!(childcontext(context), right, vn, vi)
 end
-function tilde_assume(::DefaultContext, right, vn, vi)
+function tilde_assume!!(::DefaultContext, right::Distribution, vn, vi)
     y = getindex_internal(vi, vn)
     f = from_maybe_linked_internal_transform(vi, vn, right)
     x, logjac = with_logabsdet_jacobian(f, y)
     vi = accumulate_assume!!(vi, x, logjac, vn, right)
     return x, vi
 end
-function tilde_assume(context::PrefixContext, right, vn, vi)
+function tilde_assume!!(context::PrefixContext, right::Distribution, vn, vi)
     # Note that we can't use something like this here:
     #     new_vn = prefix(context, vn)
-    #     return tilde_assume(childcontext(context), right, new_vn, vi)
+    #     return tilde_assume!!(childcontext(context), right, new_vn, vi)
     # This is because `prefix` applies _all_ prefixes in a given context to a
     # variable name. Thus, if we had two levels of nested prefixes e.g.
     # `PrefixContext{:a}(PrefixContext{:b}(DefaultContext()))`, then the
     # first call would apply the prefix `a.b._`, and the recursive call
     # would apply the prefix `b._`, resulting in `b.a.b._`.
     # This is why we need a special function, `prefix_and_strip_contexts`.
     new_vn, new_context = prefix_and_strip_contexts(context, vn)
-    return tilde_assume(new_context, right, new_vn, vi)
+    return tilde_assume!!(new_context, right, new_vn, vi)
 end
 
 """
     tilde_assume!!(context, right, 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`.
-
-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 right isa DynamicPPL.Submodel
-        _evaluate!!(right, vi, context, vn)
-    else
-        tilde_assume(context, right, vn, vi)
-    end
+function tilde_assume!!(context, right::DynamicPPL.Submodel, vn, vi)
+    return _evaluate!!(right, vi, context, vn)
 end
 
 # observe

src/contexts.jl

@@ -185,7 +185,7 @@ PrefixContexts removed.
 
 NOTE: This does _not_ modify any variables in any `ConditionContext` and
 `FixedContext` that may be present in the context stack. This is because this
-function is only used in `tilde_assume`, which is lower in the tilde-pipeline
+function is only used in `tilde_assume!!`, which is lower in the tilde-pipeline
 than `contextual_isassumption` and `contextual_isfixed` (the functions which
 actually use the `ConditionContext` and `FixedContext` values). Thus, by this
 time, any `ConditionContext`s and `FixedContext`s present have already served

src/contexts/init.jl

@@ -134,7 +134,7 @@ struct InitContext{R<:Random.AbstractRNG,S<:AbstractInitStrategy} <: AbstractCon
 end
 NodeTrait(::InitContext) = IsLeaf()
 
-function tilde_assume(
+function tilde_assume!!(
     ctx::InitContext, dist::Distribution, vn::VarName, vi::AbstractVarInfo
 )
     in_varinfo = haskey(vi, vn)

src/debug_utils.jl

@@ -253,7 +253,7 @@ function record_post_tilde_assume!(context::DebugContext, vn, dist, value, varin
     return nothing
 end
 
-function DynamicPPL.tilde_assume(context::DebugContext, right, vn, vi)
+function DynamicPPL.tilde_assume!!(context::DebugContext, right, vn, vi)
     record_pre_tilde_assume!(context, vn, right, vi)
     value, vi = DynamicPPL.tilde_assume!!(childcontext(context), right, vn, vi)
     record_post_tilde_assume!(context, vn, right, value, vi)

src/transforming.jl

@@ -12,7 +12,7 @@ how to do the transformation, used by e.g. `SimpleVarInfo`.
 struct DynamicTransformationContext{isinverse} <: AbstractContext end
 NodeTrait(::DynamicTransformationContext) = IsLeaf()
 
-function tilde_assume(
+function tilde_assume!!(
     ::DynamicTransformationContext{isinverse}, right, vn, vi
 ) where {isinverse}
     r = vi[vn, right]