Fix error in local_reshape_to_dimshuffle#2107
Merged
ricardoV94 merged 1 commit intopymc-devs:mainfrom May 1, 2026
Merged
Conversation
ricardoV94
commented
May 1, 2026
| copy_stack_trace(output, new_out) | ||
| if all(inp.type.broadcastable) or not new_output_shape: | ||
| # Trivial case we have provably size 1 as input or output, reshape can't be doing anything useful | ||
| new_out = inp.dimshuffle(["x"] * output.type.ndim) |
Member
Author
There was a problem hiding this comment.
This drops all pre-existing axes and expands again what's needed. Less work to write than trying to keep original length one axes and expand what's missing. Actually should check what pytensor likes to canonicalize into
Member
Author
There was a problem hiding this comment.
OTOH dropping all implicitly adds an assert that size==1 in Dimshuffle perform/dispatch, which is nice to have
jessegrabowski
approved these changes
May 1, 2026
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
2 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.
Found out when working on pymc-devs/pymc-extras#533 by @Michal-Novomestsky
The rewrite never anticipating not finding the input all squeezable but concluding the reshape was useless (e.g., only ones, or empty) from the reshape(x, shape ) -> shape analysis.
Now we have a clear branch for the case where reshape isn't needed (either input or output tells us we must have a size 1 input), and otherwise we know we always need squeeze -> reshape -> expand_dims. This also simplifies and we don't need to call the other rewrite to cleanup the case where reshape isn't needed