Split up tilde_observe!! for Distribution / Submodel

Author: penelopeysm

src/context_implementations.jl

@@ -60,9 +60,11 @@ 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, left, vn, vi)
-    right isa DynamicPPL.Submodel &&
-        throw(ArgumentError("`x ~ to_submodel(...)` is not supported when `x` is observed"))
+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)
+    throw(ArgumentError("`x ~ to_submodel(...)` is not supported when `x` is observed"))
+end

src/transforming.jl

@@ -13,7 +13,7 @@ struct DynamicTransformationContext{isinverse} <: AbstractContext end
 NodeTrait(::DynamicTransformationContext) = IsLeaf()
 
 function tilde_assume!!(
-    ::DynamicTransformationContext{isinverse}, right, vn, vi
+    ::DynamicTransformationContext{isinverse}, right::Distribution, vn, vi
 ) where {isinverse}
     # vi[vn, right] always provides the value in unlinked space.
     x = vi[vn, right]
@@ -31,7 +31,7 @@ function tilde_assume!!(
     return x, vi
 end
 
-function tilde_observe!!(::DynamicTransformationContext, right, left, vn, vi)
+function tilde_observe!!(::DynamicTransformationContext, right::Distribution, left, vn, vi)
     return tilde_observe!!(DefaultContext(), right, left, vn, vi)
 end