Repository navigation
Correctness of Libtask with submodels depends on whether submodels' functions are inlined or not #2772
Description
Consider the following models, which are entirely equivalent:
using Turing, Random
@model function inner(y, x)
y ~ Normal(x)
end
@model function nested(y)
x ~ Normal()
a ~ to_submodel(inner(y, x))
end
m1 = nested(1.0)
mean(sample(Xoshiro(468), m1, PG(10), 1000))
@model function single(y)
x ~ Normal()
y ~ Normal(x)
end
m2 = single(1.0)
mean(sample(Xoshiro(468), m2, PG(10), 1000))It turns out that (at least on my system) both of these sample correctly with PG (giving a mean x of 0.46).
If you think about it carefully though, there's no reason why the submodel one should work. The call to Libtask.produce is inside the y ~ Normal(x) statement, which is in turn inside the inner submodel. So, in order to get all the way to the produce call, we need to mark everything along the call stack with Libtask.might_produce. This is all basic Libtask functionality.
The problem is that the method inner(model, varinfo, ...) itself has never been marked as produce-able! So how does this sampling actually work correctly, then?
The answer is that if the submodel is small enough, then the method inner(model, varinfo, ...) will just be inlined.
We can demonstrate what happens when this fails to inline, by changing the output of the DynamicPPL compiler. If we were to apply this patch on the current release of DPPL (v0.39.13)
diff --git a/src/compiler.jl b/src/compiler.jl
index 84a9a485..604e0b32 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -710,6 +710,7 @@ function build_output(modeldef, linenumbernode)
# element in the returned value is always the most up-to-date `__varinfo__`.
# See the docstrings of `replace_returns` for more info.
evaluatordef[:body] = MacroTools.@q begin
+ @noinline
$(linenumbernode)
$(replace_returns(add_return_to_last_statment(modeldef[:body])))
endand rerun the PG sampling for the submodel, we get a mean x of 0.01, which is consistent with the call to produce never actually being triggered.
So essentially, we are pretty much at the mercy of the Julia compiler here. If it decides to inline the submodel, we are good. If it doesn't, we will silently give wrong results.
Activity
penelopeysm commentedon Feb 12, 2026
I don't immediately see a very clean answer to this. The most direct solution is that the user will have to mark all their models as might_produce themselves (just like how we force them to do it if the model has keyword arguments). The other option would be to make Libtask a DynamicPPL dependency, and modify the
@modelmacro to additionally emit extra code to mark all the newly defined methods with might_produce.penelopeysm commentedon Feb 12, 2026
Another possibility would be to write a DynamicPPLLibtask extension, that extends
check_modelto also check if any inner submodels are marked as might_produce (we can do this quite easily I think by checking ifright isa Submodeland callingLibtask.might_produceon the relevant method, which we can get frommethods(right.model.f)). In that case, the remedy is still the user's responsibility (they will get an error that will tell them to mark the submodel as might_produce) but at least we don't silently give wrong results (which is my nightmare scenario).I might try to experiment with that and see if we can make it work, as that seems like a less intrusive option than hardcoding it into the output of
@model.0AnshuAditya0 commentedon Feb 17, 2026
PR #198 in Libtask adds
@might_produce- is that directly intended to fix this? because even if Libtask exposes the macro, DynamicPPL's compiler still needs to emit it on the evaluator function for submodels. so the fix probably needs changes in both places - Libtask side to expose the mechanism, and compiler.jl to actually annotate the generated evaluator with@might_producewhenever it's building a submodel evaluator.If we go the
@inline route, we need to be careful.@inlineis technically just a hint; the compiler will ignore it if a submodel exceeds certain complexity limits or contains specific control flow. If the compiler decides not to inline a large submodel, It will failAlso the test suite concern - if all existing submodel + PG tests pass, it's probably because every test submodel happens to be small enough to inline.So tests give a false sense of security. a more complex submodel in production would silently break without anyone noticing.
penelopeysm commentedon Mar 1, 2026
@might_produceis just a thin wrapper aroundmight_produce, which has existed since the beginning of (this version of) Libtask.DynamicPPL doesn't currently depend on Libtask, so this has to be handled at the Turing level, which #2778 now does.
Mark all methods with `DynamicPPL.Model` as produceable (#2780)
penelopeysm commentedon Mar 2, 2026
Closed by #2780