Misc RandomVariable improvements #79
Merged
Conversation
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
ricardoV94
added
bug
This allows PyTensor to infer more broadcastable patterns, by placing the casting inside the MakeVector Op
ricardoV94
force-pushed
the
fix_rv_ds_lift
branch
2 times, most recently
from
a18e46a to
128531f
Compare
|
@ricardoV94 I unsubscribed - ping me when CI is green |
99.2% confident it will be green this run @michaelosthege |
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## main #79 +/- ##
==========================================
+ Coverage 74.22% 74.26% +0.03%
==========================================
Files 174 175 +1
Lines 48734 48929 +195
Branches 10367 10395 +28
==========================================
+ Hits 36175 36335 +160
- Misses 10272 10291 +19
- Partials 2287 2303 +16
|
ricardoV94
commented
Dec 3, 2022
pytensor/tensor/random/op.py
Outdated
| @@ -319,8 +319,23 @@ def make_node(self, rng, size, dtype, *dist_params): | |||
| "The type of rng should be an instance of either RandomGeneratorType or RandomStateType" | |||
| ) | |||
|
|
|||
| # Fail early when size is incompatible with parameters | |||
| size_len = get_vector_length(size) | |||
There was a problem hiding this comment.
Perhaps I should move this to infer_shape. If for some reason an RV needs different batching semantics it can override it?
There was a problem hiding this comment.
can you come up with an example?
one of the zero-sum RVs maybe?
if not I'd say apply the YAGNI rule
There was a problem hiding this comment.
There was already a case with the ChoiceRV and PermutationRVs which need special handling for the shape
michaelosthege
approved these changes
Dec 3, 2022
pytensor/tensor/random/op.py
Outdated
| @@ -319,8 +319,23 @@ def make_node(self, rng, size, dtype, *dist_params): | |||
| "The type of rng should be an instance of either RandomGeneratorType or RandomStateType" | |||
| ) | |||
|
|
|||
| # Fail early when size is incompatible with parameters | |||
| size_len = get_vector_length(size) | |||
There was a problem hiding this comment.
can you come up with an example?
one of the zero-sum RVs maybe?
if not I'd say apply the YAGNI rule
… parameters dimensionality
* The rewrite no longer bails out when dimshuffle affects both unique param dimensions and repeated param dimensions from the size argument. This requires: 1) Adding broadcastable dimensions to the parameters, which should be "cost-free" and would need to be done in the `perform` method anyway. 2) Extend size to incorporate implicit batch dimensions coming from the parameters. This requires computing the shape resulting from broadcasting the parameters. It's unclear whether this is less performant, because the `perform` method can now simply broadcast each parameter to the size, instead of having to broadcast the parameters together. * The rewrite now works with Multivariate RVs * The rewrite bails out when dimensions are dropped by the Dimshuffle. This case was not correctly handled by the previous rewrite
ricardoV94
force-pushed
the
fix_rv_ds_lift
branch
from
128531f to
e59ddb7
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
This PR improves a couple of edge issues related to RandomVariable static shape inference, and simplifies and extends the
local_dimshuffle_rv_lift rewrite.Allowing the rewrite to apply in more cases increases the range of graphs we can infer the logprob in PyMC: https://github.com/pymc-devs/pymc/blob/a0d6ba079eac2f044ed40cc5747f1079d99f9f16/pymc/logprob/tensor.py#L288-L290
Closes #60 and makes progress related to #49