Implement dims CustomDist (alternative impl)#8346
Draft
ricardoV94 wants to merge 3 commits into
Draft
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #8346 +/- ##
==========================================
+ Coverage 91.72% 91.80% +0.07%
==========================================
Files 125 129 +4
Lines 20526 20855 +329
==========================================
+ Hits 18828 19146 +318
- Misses 1698 1709 +11
🚀 New features to boost your workflow:
|
ricardoV94
force-pushed
the
alt-customdist-dims
branch
from
July 2, 2026 20:44
971bd5b to
7a470a4
Compare
Normalize size and convert params before the type check, and rebuild the Op when the param types (not just the size type) don't match the inner graph.
ricardoV94
force-pushed
the
alt-customdist-dims
branch
2 times, most recently
from
July 2, 2026 21:52
1c8004a to
ecb679b
Compare
…ions The xtensor counterpart of SymbolicRandomVariable: an OpFromGraph that dims distributions dispatch logp/logcdf/support_point on, with a single piped RNG and a generic rebuild used to add extra dims. Censored is the first user.
The dist function receives the params and extra_dims (the counterpart of size) as XTensorVariables and returns an xtensor random graph, demarcated by a DimCustomDistRV Op that logp/logcdf/support_point overrides are dispatched on with xtensor semantics. Co-authored-by: Will Dean <wd60622@gmail.com>
ricardoV94
force-pushed
the
alt-customdist-dims
branch
from
July 2, 2026 21:59
ecb679b to
978eba4
Compare
ricardoV94
commented
Jul 3, 2026
| assert c3_dist.dims == ("d", "c", "a", "b") | ||
|
|
||
|
|
||
| def test_censored_custom_dist(): |
Member
Author
There was a problem hiding this comment.
This is really test custom dist not censored, and local imports -.-
ricardoV94
commented
Jul 5, 2026
| }, | ||
| ) | ||
|
|
||
| # ---- Dispatch the user overrides with xtensor semantics ---- |
Member
Author
There was a problem hiding this comment.
Suggested change
| # ---- Dispatch the user overrides with xtensor semantics ---- |
ricardoV94
commented
Jul 5, 2026
| ) | ||
|
|
||
|
|
||
| class TestCustomDistSymbolic: |
Member
Author
There was a problem hiding this comment.
don't need this test class, there are no two forms of dist like in regular CustomDist
ricardoV94
commented
Jul 5, 2026
| logp_val = model.compile_logp()(ip) | ||
| assert np.isfinite(logp_val) | ||
|
|
||
| def test_mu_as_model_var(self): |
Member
Author
There was a problem hiding this comment.
why is this testing with tensor logp? also way too many tests
Member
Author
|
Still cleanup but feel free to take a look @williambdean Not sure I landed very far from your proposal, but this is more or less what I imagined it may look like |
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
1 participant
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.
Exploring alternative approach to #8311
Add a DimSymbolicRV analogous to SymbolicRV which we'll need for symbolic dims RV built out of other dims RV and primitives. Used by Censored and then by the new dims.CustomDist when dist is provided