Skip to content

feat: inner_dist peel hook — collapse leaf-wrapper forwarding boilerplate - #362

Closed
seabbs-bot wants to merge 2 commits into
mainfrom
feat/wrapper-inner-leaf
Closed

feat: inner_dist peel hook — collapse leaf-wrapper forwarding boilerplate#362
seabbs-bot wants to merge 2 commits into
mainfrom
feat/wrapper-inner-leaf

Conversation

@seabbs-bot

@seabbs-bot seabbs-bot commented Aug 7, 2026

Copy link
Copy Markdown
Collaborator

Adds a single-layer peel hook inner_dist(leaf) whose base identity returns the leaf itself. The read-through leaf-wrapper hooks (free_leaf, uncertain_specs, extra_leaf_params, shared_tag) now recurse through inner_dist, so a wrapper type registers ONE method returning its inner distribution instead of one redundant forwarding method per hook.

What changes

  • inner_dist(leaf) = leaf base (a plain leaf's inner is itself — the terminal) plus inner_dist(::Truncated), ::Shared, ::Uncertain, ::Varying.
  • Generic @inline recursive peels for free_leaf / uncertain_specs / extra_leaf_params / shared_tag.
  • Dropped the now-redundant per-wrapper forward methods (free_leaf(d::Shared) etc.).
  • Rebuild hooks (rewrap_leaf, set_extra_leaf_params, instantiate) and genuinely special-cased hooks (has_varying, Shared's shared_tag, Uncertain's uncertain_specs) stay explicit — they are not pure forwards.

The payoff (ModifiedDistributions extension)

Previously a leaf wrapper needed ~16 one-line peel methods. Now it needs 4 registrations:

inner_dist(d::Affine) = get_dist(d)
inner_dist(d::Weighted) = get_dist(d)
inner_dist(d::Transformed) = get_dist(d)
inner_dist(d::Modified) = get_dist(d)

and every read-through hook forwards through them automatically.

Safety

  • Backward-compatible: existing explicit wrapper methods (CensoredDistributions, the MD extension) still win by dispatch specificity; no breakage.
  • Plain-leaf hot path: inner_dist and the === check are inferred away for a concrete non-wrapper leaf, so free_leaf(::Gamma) compiles to an identity. Verified by running the package: uncertain_specs/shared_tag specific methods still win, nested wrappers peel correctly, plain leaves are identity.
  • The @generated flat codec (codec_gen.jl) uses the type-level registry, not these hooks — unaffected.
  • New @testitem covers plain/shared/truncated/nested behavior.

Note: the pre-existing 15 failures + 3 errors in the non-AD suite are in ConvolvedDistributions interop tests (convolve_series auto-discretisation #226, quantile_by_optimization #112/#337) caused by the rev=main ConvolvedDistributions pin — addressed separately by PR #359. None touch this change.

This was opened by a bot. Please ping @seabbs for any questions.

Refs #332. Unblocks EpiAware/ModifiedDistributions.jl#144.

Add a single-layer peel hook inner_dist(leaf) whose base identity returns the
leaf itself. Every read-through leaf-wrapper hook (free_leaf, uncertain_specs,
extra_leaf_params, shared_tag) now recurses through inner_dist, so a wrapper
type registers ONE method returning its inner distribution instead of one
redundant forwarding method per hook. Register inner_dist for Truncated,
Shared, Uncertain and Varying, and drop their now-redundant explicit forwards.

Rather than the 16 one-line peel methods a wrapper currently needs in the
ModifiedDistributions extension, it now needs 4 inner_dist registrations:
    inner_dist(d::Affine) = get_dist(d)
    inner_dist(d::Weighted) = get_dist(d)
    inner_dist(d::Transformed) = get_dist(d)
    inner_dist(d::Modified) = get_dist(d)
The generic peel recurses until the base identity (inner === leaf) stops it,
which is the terminal for a plain leaf. Rebuild hooks (rewrap_leaf,
set_extra_leaf_params, instantiate) and special-cased hooks (has_varying,
Shared's shared_tag, Uncertain's uncertain_specs) stay explicit where their
semantics are not a pure forward. Backward-compatible: existing explicit
wrapper methods still win by dispatch specificity over the new generics.

Closes the leaf-wrapper boilerplate half of the leaf-protocol simplification.
@github-actions

github-actions Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Try this Pull Request!

Option 1: Julia Package Manager

Open Julia and type:

import Pkg
Pkg.activate(temp=true)
Pkg.add(url="https://github.com/EpiAware/ComposedDistributions.jl", rev="feat/wrapper-inner-leaf")
using ComposedDistributions

Option 2: Local Checkout

If you have the repo locally:

git checkout feat/wrapper-inner-leaf
julia --project=. -e "using Pkg; Pkg.instantiate()"

@codecov

codecov Bot commented Aug 12, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

Flag Coverage Δ
ad-enzyme-forward 0.00% <0.00%> (ø)
ad-enzyme-reverse 13.75% <94.11%> (+0.12%) ⬆️
ad-forwarddiff 12.04% <70.58%> (-0.01%) ⬇️
ad-mooncake-forward 0.00% <0.00%> (ø)
ad-mooncake-reverse 12.24% <70.58%> (-0.01%) ⬇️
ad-reversediff 12.04% <70.58%> (-0.01%) ⬇️
unit 86.92% <100.00%> (+0.05%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
src/composers/Shared.jl 63.75% <100.00%> (+1.25%) ⬆️
src/composers/Uncertain.jl 93.66% <100.00%> (-0.05%) ⬇️
src/composers/introspection.jl 83.97% <100.00%> (+0.14%) ⬆️
src/composers/varying.jl 79.13% <100.00%> (-0.18%) ⬇️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@github-actions

Copy link
Copy Markdown
Contributor

📖 Documentation preview is ready!

View the docs for this PR at: https://EpiAware.github.io/ComposedDistributions.jl/previews/PR362/

This preview will be updated automatically when you push new commits.

@github-actions

Copy link
Copy Markdown
Contributor

Benchmark comparison vs base

Minimum time per call. Buckets are PR time as a % of base, so lower is faster (🟢 faster, ⚪ within 5%, 🔴 slower). Counts of benchmarks per bucket:

Group 🟢 <50% 🟢 50–75% 🟢 75–95% ⚪ 95–105% 🔴 105–125% 🔴 125–150% 🔴 >150%
Evaluation · 1 4 7 3 · 2
ForwardDiff · · · 9 · · ·
ReverseDiff (tape) · · 3 6 · · ·
Mooncake reverse · 1 3 5 · · ·
Enzyme reverse · · 4 5 · · ·
Evaluation — 17 benchmarks (by time change)
Benchmark base PR time memory
Composition / Choose / construct 20.0 ns 30.0 ns 🔴 1.5×
Composition / Compete / construct 20.0 ns 30.0 ns 🔴 1.5×
Composition / Resolve / rand 631.0 ns 470.0 ns 🟢 0.74× ⚪ 1.0×
Composition / Compete / rand 621.0 ns 480.0 ns 🟢 0.77× ⚪ 1.0×
Composition / Choose / logpdf 49.0 ns 60.0 ns 🔴 1.22×
Composition / Compete / logccdf 320.0 ns 350.0 ns 🔴 1.09×
Composition / Sequential / rand 1.27 μs 1.16 μs 🟢 0.91× ⚪ 1.0×
Composition / Parallel / rand 1.25 μs 1.15 μs 🟢 0.92× ⚪ 1.0×
Composition / Resolve / logpdf 130.0 ns 140.0 ns 🔴 1.08×
Composition / Parallel / construct 1.41 μs 1.33 μs 🟢 0.94× ⚪ 1.0×
Composition / Sequential / construct 1.4 μs 1.35 μs ⚪ 0.96× ⚪ 1.0×
Composition / Nested / logpdf 1.69 μs 1.63 μs ⚪ 0.96× ⚪ 1.0×
Composition / Resolve / construct 29.0 ns 30.0 ns ⚪ 1.03×
Composition / Nested / rand 4.48 μs 4.34 μs ⚪ 0.97× ⚪ 1.0×
Composition / Parallel / logpdf 100.0 ns 101.0 ns ⚪ 1.01× ⚪ 1.0×
Composition / Nested / compose 2.64 μs 2.62 μs ⚪ 1.0× ⚪ 1.0×
Composition / Sequential / logpdf 100.0 ns 100.0 ns ⚪ 1.0× ⚪ 1.0×
AD gradients — 36 benchmarks (by time change)
Benchmark base PR time memory
AD gradients / Censored leaf marginal logpdf / Mooncake reverse 429.7 μs 321.76 μs 🟢 0.75× ⚪ 1.0×
AD gradients / Sequential Gamma+LogNormal logpdf / Mooncake reverse 403.62 μs 315.49 μs 🟢 0.78× ⚪ 1.0×
AD gradients / Shared-tag unflatten/update codec / Mooncake reverse 114.18 μs 95.57 μs 🟢 0.84× 🔴 1.08×
AD gradients / Censored leaf marginal logpdf / Enzyme reverse 130.07 μs 119.84 μs 🟢 0.92× ⚪ 1.0×
AD gradients / Choose selected-branch logpdf / ReverseDiff (tape) 14.42 μs 13.39 μs 🟢 0.93× ⚪ 1.0×
AD gradients / Pool non-centred reconstruction logpdf / ReverseDiff (tape) 26.67 μs 24.78 μs 🟢 0.93× ⚪ 1.0×
AD gradients / Resolve stick-breaking branch-prob logpdf / Mooncake reverse 97.91 μs 91.12 μs 🟢 0.93× ⚪ 1.0×
AD gradients / Sequential Gamma+LogNormal logpdf / ReverseDiff (tape) 28.47 μs 26.65 μs 🟢 0.94× ⚪ 1.0×
AD gradients / Truncated uncertain-leaf unflatten/update codec / Enzyme reverse 22.78 μs 21.34 μs 🟢 0.94× ⚪ 1.0×
AD gradients / Choose selected-branch logpdf / Enzyme reverse 3.41 μs 3.19 μs 🟢 0.94× ⚪ 1.0×
AD gradients / Pool non-centred reconstruction logpdf / Enzyme reverse 2.01 μs 1.91 μs 🟢 0.95× ⚪ 1.0×
AD gradients / Truncated uncertain-leaf unflatten/update codec / ForwardDiff 3.99 μs 4.19 μs ⚪ 1.05× ⚪ 1.0×
AD gradients / Resolve stick-breaking branch-prob logpdf / ReverseDiff (tape) 55.01 μs 52.58 μs ⚪ 0.96× ⚪ 1.0×
AD gradients / Resolve mixture marginal logpdf / Mooncake reverse 97.19 μs 93.03 μs ⚪ 0.96× ⚪ 1.0×
AD gradients / Shared-tag unflatten/update codec / ReverseDiff (tape) 11.96 μs 11.46 μs ⚪ 0.96× ⚪ 1.0×
AD gradients / Censored leaf marginal logpdf / ReverseDiff (tape) 28.02 μs 26.85 μs ⚪ 0.96× ⚪ 1.0×
AD gradients / Compete racing-hazard marginal logpdf / ReverseDiff (tape) 50.0 μs 47.99 μs ⚪ 0.96× ⚪ 1.0×
AD gradients / Compete racing-hazard marginal logpdf / Mooncake reverse 52.98 μs 51.15 μs ⚪ 0.97× ⚪ 1.0×
AD gradients / Resolve stick-breaking branch-prob logpdf / ForwardDiff 6.94 μs 6.71 μs ⚪ 0.97× ⚪ 1.0×
AD gradients / Pool non-centred reconstruction logpdf / ForwardDiff 1.0 μs 972.0 ns ⚪ 0.97× ⚪ 1.0×
AD gradients / Truncated uncertain-leaf unflatten/update codec / ReverseDiff (tape) 12.29 μs 11.94 μs ⚪ 0.97× ⚪ 1.02×
AD gradients / Resolve mixture marginal logpdf / ReverseDiff (tape) 63.5 μs 61.72 μs ⚪ 0.97× ⚪ 1.0×
AD gradients / Sequential Gamma+LogNormal logpdf / ForwardDiff 10.42 μs 10.69 μs ⚪ 1.03× ⚪ 1.0×
AD gradients / Resolve stick-breaking branch-prob logpdf / Enzyme reverse 21.18 μs 20.74 μs ⚪ 0.98× ⚪ 1.0×
AD gradients / Choose selected-branch logpdf / Mooncake reverse 16.73 μs 17.07 μs ⚪ 1.02× ⚪ 1.0×
AD gradients / Sequential Gamma+LogNormal logpdf / Enzyme reverse 104.44 μs 102.65 μs ⚪ 0.98× ⚪ 1.0×
AD gradients / Choose selected-branch logpdf / ForwardDiff 701.0 ns 691.0 ns ⚪ 0.99× ⚪ 1.0×
AD gradients / Compete racing-hazard marginal logpdf / Enzyme reverse 9.98 μs 9.88 μs ⚪ 0.99× ⚪ 1.0×
AD gradients / Shared-tag unflatten/update codec / Enzyme reverse 8.12 μs 8.04 μs ⚪ 0.99× ⚪ 1.0×
AD gradients / Resolve mixture marginal logpdf / Enzyme reverse 20.87 μs 20.69 μs ⚪ 0.99× ⚪ 1.0×
AD gradients / Pool non-centred reconstruction logpdf / Mooncake reverse 21.19 μs 21.04 μs ⚪ 0.99× ⚪ 1.0×
AD gradients / Compete racing-hazard marginal logpdf / ForwardDiff 7.7 μs 7.75 μs ⚪ 1.01× ⚪ 1.0×
AD gradients / Censored leaf marginal logpdf / ForwardDiff 10.9 μs 10.84 μs ⚪ 0.99× ⚪ 1.0×
AD gradients / Truncated uncertain-leaf unflatten/update codec / Mooncake reverse 180.28 μs 179.44 μs ⚪ 1.0× 🔴 1.05×
AD gradients / Resolve mixture marginal logpdf / ForwardDiff 6.68 μs 6.66 μs ⚪ 1.0× ⚪ 1.0×
AD gradients / Shared-tag unflatten/update codec / ForwardDiff 1.8 μs 1.8 μs ⚪ 1.0× ⚪ 1.0×

@seabbs-bot

Copy link
Copy Markdown
Collaborator Author

Superseded by #381, which folds this and seven other PRs into one branch off main so the 0.2.0 contract reviews as a single diff. Its gates: test-fast 2379/2379, test-quality 255/255, full docs build clean.

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.

@seabbs-bot seabbs-bot closed this Aug 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants