Handle vectorization of Alloc nodes#1943
Merged
ricardoV94 merged 1 commit intopymc-devs:mainfrom Mar 6, 2026
Merged
Conversation
ricardoV94
force-pushed
the
alloc_vectorize
branch
from
4fd0eca to
ac9277c
Compare
jessegrabowski
approved these changes
Mar 6, 2026
| old_val, *old_shape = node.inputs | ||
| [old_alloc] = node.outputs | ||
|
|
||
| assert len(shape) == len(old_shape), ( |
Member
There was a problem hiding this comment.
Doesn't the number of dims have to change for a vectorize 🤔
Member
Author
There was a problem hiding this comment.
Alloc takes each dimension length as an input, there's no way to make a new input out of thin air with vectorization, at most you can have sone of the dimension lengths come with batch dims
| ) | ||
|
|
||
| if not all(all(s.broadcastable) for s in shape): | ||
| # May imply a non-square Alloc |
Member
There was a problem hiding this comment.
what does "non-square" alloc mean
Member
Author
There was a problem hiding this comment.
n = lscalar()
out = alloc(1.0 ,n)
batch_n = lvector(shape=(2,))
batch_out = vectorize_graph(out, {n: batch_n})
batch_out.eval({batch_n: [3, 6]})
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.
This started being a problem in the marginalRV logp for backends that hate shapes (guess who), reported by @jessegrabowski
Incredibly we had managed all this time without a specific handling other than the useless Blockwise