Skip to content

fix(bell): reindex censored_mask to leaf order in dynamic-censoring path - #10

Merged
thisiscam merged 1 commit into
mainfrom
fix-nested-censoring-mask
May 27, 2026
Merged

thisiscam merged 1 commit into
mainfrom
fix-nested-censoring-mask

Conversation

@thisiscam

Copy link
Copy Markdown
Owner

optimize_order() permutes leaves by uncensored-leaf count, but the dynamic-censoring path indexed the caller-supplied censored_mask by leaf position while the mask arrives in obs (dimension) order. Nested models with heterogeneous panel sizes AND a censored leaf therefore flagged the WRONG leaves as censored, producing incorrect nested log-densities. Homogeneous panels were unaffected (optimize_order is identity there), which is why this escaped earlier tests.

Add _mask_to_leaf_order() and reindex the mask at both Bell entry points, mirroring the obs gather in _marginal_transforms_chunked.

Validated against an independent mpmath ground truth. Adds tests/test_nested_censoring_heteropanel.py (panels [3,4]/[2,4]/[2,2,3] + uncensored control) — fails before, passes after (5/5).

optimize_order() permutes leaves by uncensored-leaf count, but the
dynamic-censoring path indexed the caller-supplied censored_mask by leaf
position while the mask arrives in obs (dimension) order. Nested models with
heterogeneous panel sizes AND a censored leaf therefore flagged the WRONG
leaves as censored, producing incorrect nested log-densities. Homogeneous
panels were unaffected (optimize_order is identity there), which is why this
escaped earlier tests.

Add _mask_to_leaf_order() and reindex the mask at both Bell entry points,
mirroring the obs gather in _marginal_transforms_chunked.

Validated against an independent mpmath ground truth. Adds
tests/test_nested_censoring_heteropanel.py (panels [3,4]/[2,4]/[2,2,3] +
uncensored control) — fails before, passes after (5/5).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
@thisiscam
thisiscam merged commit f061d00 into main May 27, 2026
2 checks passed
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.

1 participant