Remove internals.product_n and speed up fvsof and sizesof - #747
Merged
Conversation
Contributor
Author
|
I think this should have a nice side effect of speeding up the NumPyro CI build, which is the slowest part of our test suite that bottlenecks CI times. |
jfeser
requested changes
Aug 6, 2026
jfeser
left a comment
jfeser
left a comment
Contributor
There was a problem hiding this comment.
fvsof shouldn't return renamed bound operations from collection binders.
eb8680
force-pushed
the
eb-remove-productn
branch
from
August 6, 2026 15:01
7f5dafc to
988910c
Compare
eb8680
force-pushed
the
eb-remove-productn
branch
from
August 6, 2026 17:31
3840d13 to
e425d1f
Compare
Contributor
Author
|
Should be fixed now |
jfeser
self-requested a review
August 6, 2026 17:34
jfeser
approved these changes
Aug 6, 2026
jfeser
force-pushed
the
eb-remove-productn
branch
from
August 10, 2026 16:39
e425d1f to
ab9a76b
Compare
jfeser
force-pushed
the
eb-remove-productn
branch
from
August 10, 2026 16:41
ab9a76b to
e81a1dd
Compare
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.
Blocked by #745
This PR refactors
fvsof,handlers.torch.sizesofandhandlers.jax.sizesofto make use of the new caching mechanism in #745 in lieu ofproductN, which seems to speed them all up by a constant factor of ~5x without affecting their semantics. It also dropsproduct_n.pyandtest_internals_product_n.py, which are now unused. It also shares the new implementation ofsizesofacross the Torch and JAX backends, improving correctness and maintainability.