Fix type-inference failure for nested submodels#1427
Merged
Conversation
Evaluating a submodel recursed through the shared `_evaluate!!(::Model, ::AbstractVarInfo)` method. Each level of submodel nesting adds another context layer to the `Model` type, so this tripped Julia's type-inference recursion limit (`limit_type_size`): from the first level of nesting onwards the model's return type was inferred as `Any`, evaluation fell back to runtime dispatch, and both primal and gradient evaluation slowed down (~15x primal, growing with nesting depth). Submodel evaluation now calls the generated model function directly instead of going through `_evaluate!!(::Model, ...)`. This keeps the growing context type out of the recursion's dispatch path (`model.f` is a distinct function type per `@model`), so inference stays type stable for nested submodels. Adds type-stability (`@inferred`) and allocation (`iszero(allocs)`) regression tests at both the evaluation and `LogDensityFunction` levels, covering several nesting depths under `UnlinkAll`/`LinkAll`. The tests were verified to fail on the unfixed code. Also bumps the version to 0.42.1 and adds a HISTORY.md entry. See TuringLang/Turing.jl#2844
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1427 +/- ##
==========================================
+ Coverage 81.58% 81.64% +0.06%
==========================================
Files 50 50
Lines 3578 3579 +1
==========================================
+ Hits 2919 2922 +3
+ Misses 659 657 -2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
|
DynamicPPL.jl documentation for PR #1427 is available at: |
Contributor
Benchmarks @ 9a0dd67Performance Ratio: gradient time divided by log-density time. For very small models these ratios are noisy across runs and machines; raw primal and gradient timings are more reliable. The benchmarks are aimed at DynamicPPL developers and mainly catch obvious allocation or type-stability regressions. See benchmark notes for details. Main @ b310eecEnvironmentJulia Version 1.11.9 Commit 53a02c0720c (2026-02-06 00:27 UTC) Build Info: Official https://julialang.org/ release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 7763 64-Core Processor WORD_SIZE: 64 LLVM: libLLVM-16.0.6 (ORCJIT, znver3) Threads: 1 default, 0 interactive, 1 GC (on 4 virtual cores) |
Member
|
Conceptually this lgtm. I tried forcing specialisation on the function in |
Member
Author
|
it is a bit baffling, I don't think the PR description was totally clear (now updated). I did a bit of digging with codex and tell it to produce a writing, I think it's passable to be read by human, so I'll just copy and paste: Nested Submodel Inference Note
The IssueThe problem is not that The problem is the extra recursive call through _evaluate!!(::Submodel, vi, parent_context, left_vn)
-> _evaluate!!(::Model, vi)At that call site, Julia infers an argument tuple type for Model{..., PrefixContext{c, InitContext{...}}}
Model{..., PrefixContext{c, PrefixContext{b, InitContext{...}}}}Julia sees the same Tuple{typeof(_evaluate!!), Model, OnlyAccsVarInfo{...}}Once Main vs Branch Call ChainThe old path was: The fixed path is: The fix removes the nested Why The Submodel Recursion Is FineFor a submodel tilde statement, the caller already has the concrete context and submodel tilde_assume!!(context, submodel, vn, NoTemplate, vi)
-> _evaluate!!(submodel, vi, context, vn)The callee signature mostly reorders values that were already present in the caller elseif is_derived_type_from_any(unwrap_unionall(t), sources, depth)
return false # t isn't something new
endThat is why a recursive Why The Model Recursion Is Not FineThe old implementation did this inside model = contextualize(submodel.model, eval_context)
return _evaluate!!(model, vi)That constructs a fresh So Julia compares recursive Tuple{typeof(_evaluate!!), Model{..., PrefixContext{c, InitContext{...}}}, VI}
Tuple{typeof(_evaluate!!), Model{..., PrefixContext{c, PrefixContext{b, InitContext{...}}}}, VI}The second call is more complex because the Why Manual Inlining WorksThe body of args, kwargs = make_evaluate_args_and_kwargs(model, vi)
return model.f(args...; kwargs...)Doing those two lines directly inside An probe code (by codex)
using InteractiveUtils
const CC = Core.Compiler
using DynamicPPL
using Distributions
module LimitProbe
using InteractiveUtils
using DynamicPPL
const CC = Core.Compiler
const CACHE_TOKEN = gensym(:LimitProbeCacheToken)
mutable struct ProbeLog
records::Vector{NamedTuple}
remarks::Vector{NamedTuple}
end
struct ProbeInterp <: CC.AbstractInterpreter
log::ProbeLog
world::UInt
inf_params::CC.InferenceParams
opt_params::CC.OptimizationParams
inf_cache::Vector{CC.InferenceResult}
end
function ProbeInterp(log::ProbeLog=ProbeLog(NamedTuple[], NamedTuple[]);
world::UInt=Base.get_world_counter(),
inf_params::CC.InferenceParams=CC.InferenceParams(),
opt_params::CC.OptimizationParams=CC.OptimizationParams(),
)
return ProbeInterp(log, world, inf_params, opt_params, CC.InferenceResult[])
end
CC.InferenceParams(interp::ProbeInterp) = interp.inf_params
CC.OptimizationParams(interp::ProbeInterp) = interp.opt_params
CC.get_inference_world(interp::ProbeInterp) = interp.world
CC.get_inference_cache(interp::ProbeInterp) = interp.inf_cache
CC.cache_owner(::ProbeInterp) = CACHE_TOKEN
function CC.add_remark!(interp::ProbeInterp, sv::CC.InferenceState, remark)
push!(
interp.log.remarks,
(;
remark=String(remark),
caller=method_label(CC.frame_instance(sv).def),
caller_spec=compact_type(CC.frame_instance(sv).specTypes),
pc=sv.currpc,
),
)
return nothing
end
CC.add_remark!(interp::ProbeInterp, sv::CC.IRInterpretationState, remark) = return nothing
function compact_type(@nospecialize(x); limit::Int=420)
s = sprint(show, x; context=:compact => true)
return lastindex(s) <= limit ? s : string(first(s, limit), " ...<", lastindex(s), " chars>")
end
method_label(m::Method) = string(m.module, ".", m.name)
function maybe_type_name(@nospecialize(t))
u = CC.unwrap_unionall(t)
return u isa DataType ? string(u.name.module, ".", u.name.name) : string(typeof(u))
end
function arg_summary(@nospecialize(sig))
st = CC.unwrap_unionall(sig)
st isa DataType || return String[]
return [compact_type(p; limit=260) for p in st.parameters]
end
function complexity_breakdown(
@nospecialize(sig),
@nospecialize(comparison),
source_svec::Core.SimpleVector,
allowed_depth::Int,
)
st = CC.unwrap_unionall(sig)
ct = CC.unwrap_unionall(comparison)
arg_complex = NamedTuple[]
model_param_complex = NamedTuple[]
st isa DataType || return (; arg_complex, model_param_complex)
ct isa DataType || return (; arg_complex, model_param_complex)
if st.name === Tuple.name && ct.name === Tuple.name
tP = st.parameters
cP = ct.parameters
ntail = length(cP) - length(tP)
tupledepth = allowed_depth > 0 ? allowed_depth - 1 : 0
if ntail >= 0
for i in eachindex(tP)
cPi = cP[i + ntail]
tPi = tP[i]
push!(
arg_complex,
(;
index=i,
more_complex=CC.type_more_complex(
tPi, cPi, source_svec, 2, tupledepth, 0
),
sig=compact_type(tPi; limit=220),
comparison=compact_type(cPi; limit=220),
),
)
end
end
if length(tP) >= 2 && length(cP) >= 2
mt = CC.unwrap_unionall(tP[2])
mc = CC.unwrap_unionall(cP[2])
if mt isa DataType && mc isa DataType && mt.name === mc.name
mtP = mt.parameters
mcP = mc.parameters
for i in 1:min(length(mtP), length(mcP))
push!(
model_param_complex,
(;
index=i,
more_complex=CC.type_more_complex(
mtP[i], mcP[i], source_svec, 3, 0, 0
),
sig=compact_type(mtP[i]; limit=220),
comparison=compact_type(mcP[i]; limit=220),
),
)
end
end
end
end
return (; arg_complex, model_param_complex)
end
function evaluate_call_kind(method::Method, @nospecialize(sig))
method.name === :_evaluate!! || return false
method.module === DynamicPPL || return false
st = CC.unwrap_unionall(sig)
st isa DataType || return false
length(st.parameters) >= 2 || return false
st.parameters[2] <: DynamicPPL.Submodel && return :submodel
st.parameters[2] <: DynamicPPL.Model && return :model
return false
end
function edge_match_details(
interp::ProbeInterp,
frame::Union{CC.InferenceState,CC.IRInterpretationState},
method::Method,
@nospecialize(sig),
sparams::Core.SimpleVector,
hardlimit::Bool,
sv::Union{CC.InferenceState,CC.IRInterpretationState},
)
world = CC.get_inference_world(interp)
callee_method2 = CC.method_for_inference_heuristics(method, sig, sparams, world)
inf_method2 = CC.method_for_inference_limit_heuristics(frame)
token_match = callee_method2 === inf_method2
cache_owner_same = !(frame isa CC.InferenceState) ||
CC.cache_owner(frame.interp) === CC.cache_owner(interp)
cycle_any = missing
parent_present = missing
parent_cached = missing
parent_has_parent = missing
parent_usable = missing
parent_matches = missing
parent_label = missing
parent_spec = missing
cycle_parent_ok = missing
if token_match && cache_owner_same &&
(!hardlimit || CC.InferenceParams(interp).ignore_recursion_hardlimit)
cycle_any = any(p::Union{CC.InferenceState,CC.IRInterpretationState}->CC.matches_sv(p, sv), CC.callers_in_cycle(frame))
if cycle_any
cycle_parent_ok = true
else
parent = CC.cycle_parent(frame)
parent_present = parent !== nothing
if parent_present
parent_cached = CC.is_cached(parent)
parent_has_parent = CC.frame_parent(parent) !== nothing
parent_usable = parent_cached || parent_has_parent
parent_matches = parent_usable && CC.matches_sv(parent, sv)
parent_label = method_label(CC.frame_instance(parent).def)
parent_spec = compact_type(CC.frame_instance(parent).specTypes)
cycle_parent_ok = parent_matches
else
parent_cached = false
parent_has_parent = false
parent_usable = false
parent_matches = false
cycle_parent_ok = false
end
end
end
return (;
frame=method_label(CC.frame_instance(frame).def),
frame_spec=compact_type(CC.frame_instance(frame).specTypes),
token_match,
callee_token=callee_method2 === nothing ? "nothing" : method_label(callee_method2),
frame_token=inf_method2 === nothing ? "nothing" : method_label(inf_method2),
cache_owner_same,
cycle_any,
parent_present,
parent_cached,
parent_has_parent,
parent_usable,
parent_matches,
parent_label,
parent_spec,
cycle_parent_ok,
edge_matches=CC.edge_matches_sv(interp, frame, method, sig, sparams, hardlimit, sv),
)
end
function diagnose_limit_pipeline(
interp::ProbeInterp,
method::Method,
@nospecialize(sig),
sparams::Core.SimpleVector,
hardlimit::Bool,
si::CC.StmtInfo,
sv::Union{CC.InferenceState,CC.IRInterpretationState},
)
sigtuple = CC.unwrap_unionall(sig)
if !(sigtuple isa DataType)
return (; invalid_sig=true)
end
edgecycle = false
topmost = nothing
self_recursion_same_sig = false
candidates = NamedTuple[]
stack = String[]
for sv′ in CC.AbsIntStackUnwind(sv)
push!(stack, method_label(CC.frame_instance(sv′).def))
infmi = CC.frame_instance(sv′)
if method === infmi.def
same_sig = (infmi.specTypes::Type == sig::Type)
if same_sig
topmost = nothing
edgecycle = true
self_recursion_same_sig = true
push!(
candidates,
(;
frame=method_label(infmi.def),
frame_spec=compact_type(infmi.specTypes),
same_sig=true,
token_match=missing,
callee_token=missing,
frame_token=missing,
cache_owner_same=missing,
cycle_any=missing,
parent_present=missing,
parent_cached=missing,
parent_has_parent=missing,
parent_usable=missing,
parent_matches=missing,
parent_label=missing,
parent_spec=missing,
cycle_parent_ok=missing,
edge_matches=missing,
),
)
break
end
if topmost === nothing
details = edge_match_details(interp, sv′, method, sig, sparams, hardlimit, sv)
push!(candidates, merge(details, (; same_sig=false)))
if details.edge_matches
topmost = sv′
edgecycle = true
end
end
end
end
washardlimit = hardlimit
hardlimit_after = hardlimit
comparison = nothing
comparison_source = "none"
source = nothing
newsig = sig
limit_changed = false
more_complex = false
spec_len = missing
arg_complex = NamedTuple[]
model_param_complex = NamedTuple[]
if topmost !== nothing
msig = CC.unwrap_unionall(method.sig)::DataType
spec_len = length(msig.parameters) + 1
mi = CC.frame_instance(sv)
if isdefined(method, :recursion_relation)
hardlimit_after = true
end
if method === mi.def
comparison = mi.specTypes
l_comparison = length((CC.unwrap_unionall(comparison)::DataType).parameters)
spec_len = max(spec_len, l_comparison)
comparison_source = "current_frame_specTypes"
elseif !hardlimit_after && topmost isa CC.InferenceState
comparison = CC.frame_instance(topmost).specTypes
comparison_source = "topmost_specTypes"
else
comparison = method.sig
comparison_source = "method_sig"
end
source = hardlimit_after ? comparison : mi.specTypes
allowed_depth = CC.InferenceParams(interp).tuple_complexity_limit_depth
source_svec = Core.svec(CC.unwrap_unionall(comparison), CC.unwrap_unionall(source))
source_svec[1] === source_svec[2] && (source_svec = Core.svec(source_svec[1]))
more_complex = CC.type_more_complex(sig, comparison, source_svec, 1, allowed_depth, spec_len)
newsig = CC.limit_type_size(sig, comparison, source, allowed_depth, spec_len)
limit_changed = newsig !== sig
breakdown = complexity_breakdown(sig, comparison, source_svec, allowed_depth)
arg_complex = breakdown.arg_complex
model_param_complex = breakdown.model_param_complex
end
return (;
invalid_sig=false,
kind=evaluate_call_kind(method, sig),
caller=method_label(CC.frame_instance(sv).def),
caller_spec=compact_type(CC.frame_instance(sv).specTypes),
method=method_label(method),
method_sig=compact_type(method.sig),
sig=compact_type(sig),
sig_hash=UInt(hash(sig)),
args=arg_summary(sig),
hardlimit_in=washardlimit,
hardlimit_after,
topmost_found=topmost !== nothing,
edgecycle,
self_recursion_same_sig,
comparison_source,
comparison=comparison === nothing ? "nothing" : compact_type(comparison),
source=source === nothing ? "nothing" : compact_type(source),
spec_len,
more_complex,
limit_changed,
newsig=compact_type(newsig),
newsig_hash=UInt(hash(newsig)),
call_result_unused=CC.call_result_unused(si),
arg_complex,
model_param_complex,
candidates,
stack=stack[1:min(end, 24)],
)
end
function CC.abstract_call_method(
interp::ProbeInterp,
method::Method,
@nospecialize(sig),
sparams::Core.SimpleVector,
hardlimit::Bool,
si::CC.StmtInfo,
sv::Union{CC.InferenceState,CC.IRInterpretationState},
)
if evaluate_call_kind(method, sig) !== false
push!(
interp.log.records,
diagnose_limit_pipeline(interp, method, sig, sparams, hardlimit, si, sv),
)
end
return @invoke CC.abstract_call_method(
interp::CC.AbstractInterpreter,
method::Method,
sig::Any,
sparams::Core.SimpleVector,
hardlimit::Bool,
si::CC.StmtInfo,
sv::Union{CC.InferenceState,CC.IRInterpretationState},
)
end
end # module LimitProbe
using .LimitProbe
@model t2844_inner() = (x ~ Normal(); return (; x))
@model t2844_middle() = (a ~ to_submodel(t2844_inner()); return (; x=a.x))
@model t2844_outer() = (b ~ to_submodel(t2844_middle()); return (; x=b.x))
@model t2844_deeper() = (c ~ to_submodel(t2844_outer()); return (; x=c.x))
function probe_init(model, tfm)
accs = setacc!!(OnlyAccsVarInfo(), LogPriorAccumulator())
return init!!(model, accs, InitFromPrior(), tfm)
end
function print_record(io, rec, i)
println(io, "=== RECORD ", i, " ===")
for key in (
:caller,
:kind,
:method,
:sig_hash,
:newsig_hash,
:hardlimit_in,
:hardlimit_after,
:topmost_found,
:edgecycle,
:self_recursion_same_sig,
:comparison_source,
:spec_len,
:more_complex,
:limit_changed,
:call_result_unused,
)
println(io, key, " = ", getproperty(rec, key))
end
println(io, "args:")
for (j, arg) in pairs(rec.args)
println(io, " arg", j, " = ", arg)
end
println(io, "comparison = ", rec.comparison)
println(io, "source = ", rec.source)
println(io, "sig = ", rec.sig)
println(io, "newsig = ", rec.newsig)
println(io, "candidates:")
for (j, c) in pairs(rec.candidates)
println(io, " candidate ", j)
for key in (
:frame,
:same_sig,
:edge_matches,
:token_match,
:callee_token,
:frame_token,
:cache_owner_same,
:cycle_any,
:parent_present,
:parent_cached,
:parent_has_parent,
:parent_usable,
:parent_matches,
:cycle_parent_ok,
:parent_label,
)
println(io, " ", key, " = ", getproperty(c, key))
end
println(io, " frame_spec = ", c.frame_spec)
println(io, " parent_spec = ", c.parent_spec)
end
println(io, "arg_complex:")
for c in rec.arg_complex
println(
io,
" arg",
c.index,
" more_complex = ",
c.more_complex,
" :: ",
c.sig,
" vs ",
c.comparison,
)
end
println(io, "model_param_complex:")
for c in rec.model_param_complex
println(
io,
" param",
c.index,
" more_complex = ",
c.more_complex,
" :: ",
c.sig,
" vs ",
c.comparison,
)
end
println(io, "stack:")
for frame in rec.stack
println(io, " ", frame)
end
end
label = get(ENV, "DPPL_PROBE_LABEL", "unknown")
tfm = get(ENV, "DPPL_PROBE_TFM", "unlink") == "link" ? LinkAll() : UnlinkAll()
model = t2844_deeper()
log = LimitProbe.ProbeLog(NamedTuple[], NamedTuple[])
interp = LimitProbe.ProbeInterp(log)
rt = only(Base.return_types(probe_init, (typeof(model), typeof(tfm)); interp=interp))
println("LABEL = ", label)
println("VERSION = ", VERSION)
println("PROJECT = ", Base.active_project())
println("RETURN_TYPE = ", rt)
println("RECORD_COUNT = ", length(log.records))
println("REMARK_COUNT = ", length(log.remarks))
println(
"BOUNDED_REMARK_COUNT = ",
count(rem -> occursin("Bounded recursion", rem.remark), log.remarks),
)
for (i, rec) in pairs(log.records)
print_record(stdout, rec, i)
end
println("=== REMARKS ===")
for (i, rem) in pairs(log.remarks)
occursin("Bounded recursion", rem.remark) || continue
println("remark ", i, ": ", rem.remark)
println(" caller = ", rem.caller)
println(" caller_spec = ", rem.caller_spec)
println(" pc = ", rem.pc)
end |
The explanation lives in the in-code comments in src/submodel.jl and src/model.jl, so the standalone markdown note is redundant. Claude-Session: https://claude.ai/code/session_014Noe4UB3FFurwpt2BNUE5E
Keep the full mechanism explanation in src/submodel.jl; reduce the duplicated explanations in the tests and HISTORY to concise pointers to the PR. Remove the redundant NOTE block in src/model.jl. Claude-Session: https://claude.ai/code/session_014Noe4UB3FFurwpt2BNUE5E
Member
Author
|
@penelopeysm could you give this a review maybe? |
Member
|
I still don't really understand, but that's fine. I fear it might be one of those things that can't be understood without looking at the Julia source code... |
penelopeysm
approved these changes
Jun 22, 2026
Member
Author
|
yeah, I don't see it break anything and it seems to really have fixed the issue, so no real downside |
|
Thank you @sunxd3 and @penelopeysm for taking a look at this so quickly! :-) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Evaluating a submodel used to recurse through the shared
_evaluate!!(::Model, ::AbstractVarInfo)method. Each submodel level constructs a contextualizedModel{..., Ctx}whoseCtxparameter gains another context layer. In nested submodels, Julia's recursion limiter (limit_type_size) sees the same_evaluate!!(::Model, ...)method recurring with a more complexModel{..., Ctx}call signature, widens theModelargument to abstractModel, and inference can no longer resolvemodel.fprecisely. The model return type then collapses toAny, evaluation falls back to runtime dispatch, and primal/gradient evaluation slows down.Submodel evaluation now calls the generated model function directly instead of going through
_evaluate!!(::Model, ...). This removes the recursive_evaluate!!(::Model, ...)call edge while leaving the recursive_evaluate!!(::Submodel, ...)edge intact; the remaining submodel edge is not widened by the limiter. As a result, inference stays type stable for nested submodels.Adds type-stability (
@inferred) and allocation (iszero(allocs)) regression tests at both the evaluation andLogDensityFunctionlevels, covering several nesting depths underUnlinkAll/LinkAll. The tests were verified to fail on the unfixed code. Also bumps the version to 0.42.1 and adds a HISTORY.md entry.See TuringLang/Turing.jl#2844
Above is written by Claude/Codex, lightly touched by me.