Support marginalize of broadcasted normal-normal (one to many) - #715
Draft
ricardoV94 wants to merge 4 commits into
Draft
Support marginalize of broadcasted normal-normal (one to many)#715ricardoV94 wants to merge 4 commits into
ricardoV94 wants to merge 4 commits into
Conversation
ricardoV94
commented
Jul 19, 2026
ricardoV94
commented
Member
- Avoid duplicating shared RVs when unmarginalizing
- Control for broadcasting in normal-normal marginalization
- Support affine dependent means in normal-normal marginalization
- Support broadcast normal-normal marginalization as MvNormal
local_unmarginalize wraps the unmarginalized RV in a ModelFreeRV and graph_replaces it into the dependents. Because the wrapper re-references the variable being replaced, graph_replace cloned the surrounding body, duplicating any shared upstream RV wrapped in an already-built model variable (e.g. a previously unmarginalized parent). Pin those opaque ModelValuedVar wrappers as boundaries so their subgraphs are not cloned.
Register a pre-canonicalize pass ahead of the marginal rewrites (mirroring pymc.logprob) so the structure detectors see flattened Add/Mul. normal-normal now handles any affine dependent mean a + b*x (extracted by affine_coefficients), not just x or x + c. Because marginalize now returns a canonical model, equivalent_models gains an opt-in canonicalize= flag used by the round-trip test.
A latent broadcast into a wider dependent is shared across those dependents, so the marginal is a correlated MvNormal rather than the elementwise Normal. Classify each dependent dim with the dim-connection analysis (stored as dims_connections on the op): dims where the latent is broadcast become the MvNormal event (raveled and moved right), while dims matched one-to-one stay independent batch dims. The conditional mirrors this, summing each shared observation's evidence back onto the latent. The dim analysis is now the structural guard (it rejects couplings the closed form can't express, e.g. x[None,:] + x[:,None]), so the affine check accepts any DimShuffle of the latent.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #715 +/- ##
===========================================
+ Coverage 51.60% 64.66% +13.05%
===========================================
Files 73 92 +19
Lines 8003 8595 +592
===========================================
+ Hits 4130 5558 +1428
+ Misses 3873 3037 -836
🚀 New features to boost your workflow:
|
zaxtax
previously approved these changes
Jul 20, 2026
zaxtax
self-requested a review
July 20, 2026 09:24
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.