Fix conditioning in submodels

Author: penelopeysm

src/compiler.jl

@@ -53,7 +53,9 @@ function isassumption(
     vn=AbstractPPL.drop_escape(varname(expr, need_concretize(expr))),
 )
     return quote
-        if $(DynamicPPL.contextual_isassumption)(__context__, $vn)
+        if $(DynamicPPL.contextual_isassumption)(
+            __context__, $(DynamicPPL.prefix)(__context__, $vn)
+        )
             # Considered an assumption by `__context__` which means either:
             # 1. We hit the default implementation, e.g. using `DefaultContext`,
             #    which in turn means that we haven't considered if it's one of
@@ -112,8 +114,10 @@ function contextual_isassumption(context::ConditionContext, vn)
     # so we defer to `childcontext` if we haven't concluded that anything yet.
     return contextual_isassumption(childcontext(context), vn)
 end
-function contextual_isassumption(context::PrefixContext, vn)
-    return contextual_isassumption(childcontext(context), prefix(context, vn))
+function contextual_isassumption(context::PrefixContext{Prefix}, vn) where {Prefix}
+    return contextual_isassumption(
+        prefix_conditioned_variables(childcontext(context), VarName{Prefix}()), vn
+    )
 end
 
 isfixed(expr, vn) = false
@@ -473,7 +477,9 @@ function generate_tilde(left, right)
         else
             # If `vn` is not in `argnames`, we need to make sure that the variable is defined.
             if !$(DynamicPPL.inargnames)($vn, __model__)
-                $left = $(DynamicPPL.getconditioned_nested)(__context__, $vn)
+                $left = $(DynamicPPL.getconditioned_nested)(
+                    __context__, $(DynamicPPL.prefix)(__context__, $vn)
+                )
             end
 
             $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)(

src/context_implementations.jl

@@ -104,12 +104,34 @@ probability of `vi` with the returned value.
 """
 function tilde_assume!!(context, right, vn, vi)
     return if is_rhs_model(right)
-        # Prefix the variables using the `vn`.
-        rand_like!!(
-            right,
-            should_auto_prefix(right) ? PrefixContext{Symbol(vn)}(context) : context,
-            vi,
-        )
+        # We want to merge the context inside `right`, as well as the context passed
+        # down from the parent model. To ensure that _only_ variables inside
+        # right.model.context are prefixed, we need to make sure to _only_ apply
+        # PrefixContext to it, before wrapping the whole prefixed context in the
+        # parent context.
+        # 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{Symbol(vn)}(dppl_model.context)
+            # Having done this, we can set the submodel context to its leaf context.
+            new_dppl_model = contextualize(dppl_model, leafcontext(dppl_model.context))
+            # Then reconstruct the Sampleable
+            right = to_submodel(new_dppl_model, true)
+            # We can splice that prefixed submodel context into the parent context.
+            original_parent_leafcontext = leafcontext(context)
+            new_parent_context = setleafcontext(context, prefixed_submodel_context)
+            new_parent_context = setleafcontext(
+                new_parent_context, original_parent_leafcontext
+            )
+        else
+            # If we don't need to prefix, we can just use the parent context as-is
+            new_parent_context = context
+        end
+        rand_like!!(right, new_parent_context, vi)
     else
         value, logp, vi = tilde_assume(context, right, vn, vi)
         value, acclogp_assume!!(context, vi, logp)

src/contexts.jl

@@ -265,8 +265,8 @@ end
 
 Apply the prefixes in the context `ctx` to the variable name `vn`.
 """
-function prefix(ctx::PrefixContext{Prefix}, vn::VarName{Sym}) where {Prefix,Sym}
-    return AbstractPPL.prefix(prefix(childcontext(ctx), vn), VarName{Symbol(Prefix)}())
+function prefix(ctx::PrefixContext{Prefix}, vn::VarName) where {Prefix}
+    return AbstractPPL.prefix(prefix(childcontext(ctx), vn), VarName{Prefix}())
 end
 function prefix(ctx::AbstractContext, vn::VarName)
     return prefix(NodeTrait(ctx), ctx, vn)
@@ -351,6 +351,43 @@ NodeTrait(::ConditionContext) = IsParent()
 childcontext(context::ConditionContext) = context.context
 setchildcontext(parent::ConditionContext, child) = ConditionContext(parent.values, child)
 
+"""
+    prefix_conditioned_variables(context::AbstractContext, prefix::VarName)
+
+Prefix all the conditioned variables in a given context with `prefix`.
+
+```jldoctest
+julia> using DynamicPPL: prefix_conditioned_variables, ConditionContext
+
+julia> c1 = ConditionContext((a=1, ))
+ConditionContext((a = 1,), DefaultContext())
+
+julia> prefix_conditioned_variables(c1, @varname(y))
+ConditionContext(Dict(y.a => 1), DefaultContext())
+```
+"""
+function prefix_conditioned_variables(ctx::ConditionContext, prefix::VarName)
+    # Replace the prefix of the conditioned variables
+    vn_dict = to_varname_dict(ctx.values)
+    prefixed_vn_dict = Dict(
+        AbstractPPL.prefix(vn, prefix) => value for (vn, value) in vn_dict
+    )
+    # Prefix the child context as well
+    prefixed_child_ctx = prefix_conditioned_variables(childcontext(ctx), prefix)
+    return ConditionContext(prefixed_vn_dict, prefixed_child_ctx)
+end
+function prefix_conditioned_variables(c::AbstractContext, prefix::VarName)
+    return prefix_conditioned_variables(
+        NodeTrait(prefix_conditioned_variables, c), c, prefix
+    )
+end
+prefix_conditioned_variables(::IsLeaf, context::AbstractContext, prefix::VarName) = context
+function prefix_conditioned_variables(::IsParent, context::AbstractContext, prefix::VarName)
+    return setchildcontext(
+        context, prefix_conditioned_variables(childcontext(context), prefix)
+    )
+end
+
 """
     hasconditioned(context::AbstractContext, vn::VarName)
 
@@ -370,7 +407,9 @@ Return value of `vn` in `context`.
 function getconditioned(context::AbstractContext, vn::VarName)
     return error("context $(context) does not contain value for $vn")
 end
-getconditioned(context::ConditionContext, vn::VarName) = getvalue(context.values, vn)
+function getconditioned(context::ConditionContext, vn::VarName)
+    return getvalue(context.values, vn)
+end
 
 """
     hasconditioned_nested(context, vn)
@@ -387,8 +426,10 @@ hasconditioned_nested(::IsLeaf, context, vn) = hasconditioned(context, vn)
 function hasconditioned_nested(::IsParent, context, vn)
     return hasconditioned(context, vn) || hasconditioned_nested(childcontext(context), vn)
 end
-function hasconditioned_nested(context::PrefixContext, vn)
-    return hasconditioned_nested(childcontext(context), prefix(context, vn))
+function hasconditioned_nested(context::PrefixContext{Prefix}, vn) where {Prefix}
+    return hasconditioned_nested(
+        prefix_conditioned_variables(childcontext(context), VarName{Prefix}()), vn
+    )
 end
 
 """
@@ -405,8 +446,10 @@ end
 function getconditioned_nested(::IsLeaf, context, vn)
     return error("context $(context) does not contain value for $vn")
 end
-function getconditioned_nested(context::PrefixContext, vn)
-    return getconditioned_nested(childcontext(context), prefix(context, vn))
+function getconditioned_nested(context::PrefixContext{Prefix}, vn) where {Prefix}
+    return getconditioned_nested(
+        prefix_conditioned_variables(childcontext(context), VarName{Prefix}()), vn
+    )
 end
 function getconditioned_nested(::IsParent, context, vn)
     return if hasconditioned(context, vn)

src/utils.jl

@@ -1286,17 +1286,20 @@ broadcast_safe(x::Distribution) = (x,)
 broadcast_safe(x::AbstractContext) = (x,)
 
 # Convert (x=1,) to Dict(@varname(x) => 1)
-_nt_to_varname_dict(nt) = Dict(VarName{k}() => v for (k, v) in pairs(nt))
+function to_varname_dict(nt::NamedTuple)
+    return Dict{VarName,Any}(VarName{k}() => v for (k, v) in pairs(nt))
+end
+to_varname_dict(d::AbstractDict) = d
 # Version of `merge` used by `conditioned` and `fixed` to handle
 # the scenario where we might try to merge a dict with an empty
 # tuple.
 # TODO: Maybe replace the default of returning `NamedTuple` with `nothing`?
 _merge(left::NamedTuple, right::NamedTuple) = merge(left, right)
 _merge(left::AbstractDict, right::AbstractDict) = merge(left, right)
 _merge(left::AbstractDict, ::NamedTuple{()}) = left
-_merge(left::AbstractDict, right::NamedTuple) = merge(left, _nt_to_varname_dict(right))
+_merge(left::AbstractDict, right::NamedTuple) = merge(left, to_varname_dict(right))
 _merge(::NamedTuple{()}, right::AbstractDict) = right
-_merge(left::NamedTuple, right::AbstractDict) = merge(_nt_to_varname_dict(left), right)
+_merge(left::NamedTuple, right::AbstractDict) = merge(to_varname_dict(left), right)
 
 """
     unique_syms(vns::T) where {T<:NTuple{N,VarName}}