fix(pool): dispatch-drive the flatten walk, fix vacuous MomentLeaf guard (#352) - #363
Closed
seabbs-bot wants to merge 4 commits into
Closed
fix(pool): dispatch-drive the flatten walk, fix vacuous MomentLeaf guard (#352)#363seabbs-bot wants to merge 4 commits into
seabbs-bot wants to merge 4 commits into
Conversation
#352) _leaf_flatten_grouped/_leaf_flatten_walk_grouped (Pool.jl) call leaf_param_names(leaf) at runtime and recurse over the result as a plain Tuple. Every built-in leaf's param_names is a literal tuple, so this happens to constant-fold and @inferred passes today -- which is why the existing pooling and S2 codec-parity tests never caught the gap. Add a leaf whose param_names comes from a Ref read (not effect-free, so it can't fold even though its TYPE still infers concretely) inside a pooled tree. @inferred flatten fails: "return type Vector{Float64} does not match inferred return type Any". Red, on purpose -- the fix lands next.
… Tuple (#352) _leaf_flatten_grouped/_leaf_flatten_walk_grouped derived their walk order by calling leaf_param_names(leaf) at runtime and recursing over the result as a plain Tuple, hoping the whole call constant-folded so `pname` stayed a Core.Const at each step. It does, for every built-in leaf (param_names returns a literal tuple), which is why nothing caught this. A leaf whose native names come from anything less trivially foldable (a Ref, a field behind a branch -- a realistic third-party Distributions.jl leaf) breaks the fold: `specs[pname]` can no longer resolve which field of the heterogeneous spec NamedTuple it needs, and the walk -- and `flatten` after it -- silently widens to Union/Any. Same class of bug _hyper_flatten_walk was already fixed for one level down (see its comment); this was the one remaining flatten-direction walker still on the old pattern. Fix: derive `names` from `entry`'s own NamedTuple{names} type parameter (entry's keys ARE leaf_param_names(leaf) by construction -- _leaf_entry_grouped built it that way) and dispatch the whole recursion on Val{names}/Val{speckeys}/Val{pool_names}/Val{materialize}, mirroring _leaf_extract (introspection.jl) and _hyper_flatten_walk (this file). `pname` is then part of the TYPE at every step, so it is a Core.Const regardless of whether leaf_param_names itself folds. Measured: @inferred flatten(tree, nt) on a pooled tree over the adversarial leaf now infers Vector{Float64} (was: Any). Every existing pooling/codec test still passes unchanged -- this is a pure robustness fix, not a behaviour change; the built-in-leaf flat layout, order and values are untouched (confirmed byte-for-byte on the S2 parity fixtures). unflatten's dual (_leaf_entry_grouped/_leaf_walk_grouped) has the same shape and the same theoretical exposure but is out of this item's scope (#352 names only the flatten-direction walker); flagged as a follow-up.
… test (#352) `isconcretetype(typeof(nt))` can never fail -- `typeof` of any runtime value is always concrete, whether or not the CALL that produced it inferred concretely. Replace it with `@inferred unflatten(...)`, matching the S2 parity item just above, and add `@inferred` to the `flatten` call too (it had no inference assertion at all). Both pass unchanged: MomentLeaf's param_names is a literal tuple, so this item was never exposed to the runtime-name-walk gap the Pool.jl fix addresses; this is pure test hardening, not a bug fix.
Contributor
Try this Pull Request!Option 1: Julia Package ManagerOpen Julia and type: import Pkg
Pkg.activate(temp=true)
Pkg.add(url="https://github.com/EpiAware/ComposedDistributions.jl", rev="fix/352-codec-inference-hardening")
using ComposedDistributionsOption 2: Local CheckoutIf you have the repo locally: git checkout fix/352-codec-inference-hardening
julia --project=. -e "using Pkg; Pkg.instantiate()" |
Codecov Report❌ Patch coverage is
Flags with carried forward coverage won't be shown. Click here to find out more.
... and 1 file with indirect coverage changes 🚀 New features to boost your workflow:
|
Contributor
|
📖 Documentation preview is ready! View the docs for this PR at: https://EpiAware.github.io/ComposedDistributions.jl/previews/PR363/ This preview will be updated automatically when you push new commits. |
Contributor
Benchmark comparison vs baseMinimum time per call. Buckets are PR time as a % of base, so lower is faster (🟢 faster, ⚪ within 5%, 🔴 slower). Counts of benchmarks per bucket:
Evaluation — 17 benchmarks (by time change)
AD gradients — 36 benchmarks (by time change)
|
Collaborator
Author
|
Superseded by #381, which folds this and seven other PRs into one branch off This branch is preserved and the PR body is kept for its rationale — the detail here is the record of why, which #381 summarises rather than repeats. Closing so it does not compete for review. This comment was posted by a bot. Please ping @seabbs for any questions. |
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
2 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.
What changed
Two follow-ups from the S3 inference fix, per #352.
_leaf_flatten_grouped/_leaf_flatten_walk_grouped(src/composers/Pool.jl) derived their walk order by callingleaf_param_names(leaf)at runtime and recursing over the result as a plainTuple, relying on the whole call constant-folding sopnamestayed aCore.Constat each recursion step. Converted to theVal{names}-driven dispatch recursion its duals (_leaf_extractinintrospection.jl,_hyper_flatten_walkjust below in the same file) already use:namesis now taken fromentry's ownNamedTuple{names}type parameter (its keys areleaf_param_names(leaf)by construction —_leaf_entry_groupedbuilt it that way) rather than a fresh runtime call.The MomentLeaf testitem (
test/composers/codec_gen.jl, S3) carried@test isconcretetype(typeof(nt)).typeofof any runtime value is always concrete, so this assertion could never fail regardless of whether the call that producedntinferred concretely. Replaced with@inferred unflatten(...), matching the S2 parity item, and added@inferredto theflattencall too (it had no inference assertion at all) — inference is now asserted in both codec directions.Why
For every built-in leaf,
param_namesreturns a literal tuple, so_leaf_flatten_walk_grouped's constant-folding gamble happened to pay off —@inferredpassed on every existing pooling/S2 fixture, which is why nothing caught this. It only breaks for a leaf whose native parameter names come from anything the compiler can't fold to a compile-time constant (aRefread, a field behind a branch — a realistic third-partyDistributions.jlleaf that just implements the plain contract, not anything CD-specific). I reproduced this with such a leaf under apool(...)group:@code_warntype flatten(tree, nt)showedBody::ANYbefore the fix,Body::Vector{Float64}after. Oncepname's value stops being const,specs[pname](indexing a heterogeneousNamedTupleby a non-literalSymbol) can't resolve which field it needs, and the whole walk — andflattenafter it — silently widens. No error, just a boxed,Union/Any-typed result on what the Pool.jl comments already call the per-gradient hot path.This is a pure robustness fix. The flat layout, order and values for every built-in-leaf tree are unchanged — confirmed byte-for-byte on the existing S2 parity fixtures and every pooling.jl test.
Breaking surface
None. No public API, export, or docstring changed. Internal function signatures changed (
_leaf_flatten_grouped'sentryparameter is nowNamedTuple{names}instead of bareNamedTuple) but these are underscore-prefixed internals with no external callers.Scope note
unflatten's dual (_leaf_entry_grouped/_leaf_walk_grouped) has the same runtime-name-walk shape and, on inspection, the same theoretical exposure (confirmed with the same adversarial leaf —@code_warntype unflatten(...)showed an abstractNamedTuplereturn before any change here). #352 names only the flatten-direction walker, so I left it alone and flagged it as a follow-up rather than expanding scope — will file a tracking issue.Gates (this branch, Julia 1.12.6, local — CI is currently blocked on #359's
[sources]/compat fix, unrelated to this PR; I unblocked instantiate locally with the same compat bump #359 makes, without committing it)task test-fast(Pkg.test(test_args=["skip_quality"])): 1697/1697 pass, cold, twice in a row.task test-quality(Pkg.test(test_args=["quality_only"])): 239/239 pass (Aqua, ExplicitImports, doctest, formatting, JET, extension ambiguities all included).task formatcheck: clean, no changes needed after the commits.@exampletouched.task test-ad) not run — outside this item's stated gates; the fix only affectsflatten's return-type inference, not values, andreconstruct/update(the path AD differentiates through) are untouched.Commits
Red/green TDD, as requested:
test(pool): pin flatten-walk inference gap for non-foldable leaf names— adds the adversarial-leaf regression test; confirmed it fails against unfixedPool.jl(return type Vector{Float64} does not match inferred return type Any).fix(pool): drive flatten-walk recursion off Val{names}, not a runtime Tuple— the fix; confirmed green.test(codec): assert @inferred on both codec directions in the S3 leaf test— the MomentLeaf hardening (unrelated to the Pool.jl bug — MomentLeaf'sparam_namesis a literal tuple, so this was never exposed to the gap; pure test strengthening).Orthogonal to #359/#346/#362/#343 — only touches
src/composers/Pool.jland two test files, no conflict expected.This was opened by a bot. Please ping @seabbs for any questions.
Closes #352.