Conditioning on submodel variables

[penelopeysm]

This behaviour with submodels is all very reasonable:

using DynamicPPL, Distributions

@model function f()
    x ~ Normal()
    y ~ Normal()
end
@model function g()
    a ~ to_submodel(f())
end
keys(VarInfo(g()))
# 2-element Vector{VarName{sym, typeof(identity)} where sym}:
#  a.x
#  a.y

Now, let's say we wanted to condition x in the inner model. We can do that from the very outermost layer, by conditioning the model g(). When looked at from the outside, the x in the inner model is actually var"a.x", so that's what we need to use in the conditioning values. We see that this works perfectly:

cg = g() | (@varname(a.x => 1)
keys(VarInfo(cg))
# 1-element Vector{VarName{Symbol("a.y"), typeof(identity)}}:
#  a.y

Now if we instead wanted to condition the submodel itself (rather than the outermost model), one should expect that we can do that without prefixing. However, it doesn't work:

@model function h()
    a ~ to_submodel(f() | (x = 1,))
end
keys(VarInfo(h()))
# 2-element Vector{VarName{sym, typeof(identity)} where sym}:
#  a.x
#  a.y

To condition on the inner model, you still have to include the prefix:

@model function h2()
    a ~ to_submodel(f() | (@varname(a.x => 1)
end
keys(VarInfo(h2()))
# 1-element Vector{VarName{Symbol("a.y"), typeof(identity)}}:
#  a.y

This is quite counterintuitive and opens up things like this:

cf = f() | (x = 1,)
keys(VarInfo(cf)) # [y]

@model function h3()
    a ~ to_submodel(cf)
end
keys(VarInfo(h3())) # expected [a.y]; but this is [a.x, a.y]

(Note that old @submodel had the same issue.)

I didn't test fix; it might have the same problem.

[penelopeysm]

this is illustrative

using DynamicPPL, Distributions
@model function f()
    @show "inner, $(__context__)"
    x ~ Normal()
    y ~ Normal()
    return (x, y)
end
@model function g()
    @show "outer, $(__context__)"
    a ~ to_submodel(f() | (@varname(x) => 1))
end
model = g()

julia> model()
"outer, $(__context__)" = "outer, SamplingContext{SampleFromPrior, DefaultContext, Random.TaskLocalRNG}(Random.TaskLocalRNG(), SampleFromPrior(), DefaultContext())"
"inner, $(__context__)" = "inner, PrefixContext{:a, SamplingContext{SampleFromPrior, ConditionContext{Dict{VarName{:x, typeof(identity)}, Int64}, DefaultContext}, Random.TaskLocalRNG}}(SamplingContext{SampleFromPrior, ConditionContext{Dict{VarName{:x, typeof(identity)}, Int64}, DefaultContext}, Random.TaskLocalRNG}(Random.TaskLocalRNG(), SampleFromPrior(), ConditionContext(Dict(x => 1), DefaultContext())))"
(-0.6695328944144235, 0.1774662210135376)

The simplest solution might be to make to_submodel prefix all conditioned/fixed values with its prefix too (?)

DynamicPPL.jl/src/context_implementations.jl

Lines 106 to 113 in 0810e14

	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,
	)
	else

basically here we need to dig into right, get its model, get its context, and modify the values inside any ConditionContext/FixedContext as well (i don't think there's code to do that)

[penelopeysm]

Closed by #892, released in 0.36 #829