Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cleanup: toward merging core.concrete_aval & xla.abstractify #25456

Merged
merged 1 commit into from
Dec 17, 2024

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Dec 13, 2024

The goal of this change is to merge core.concrete_aval with xla.abstractify. The two functions are quite similar, but differ slightly in the details. In particular, core.concrete_aval includes sharding information while xla.abstractify does not. Additionally, xla.abstractify is in the dispatch path, and therefore is more optimized in some cases.

This first change simply moves xla.abstractify to core.abstractify, and moves the implementation of the two functions to the same locations so that they're easier to compare. This PR does not change any of the rule registrations; I plan to do that in a followup in order to localize potentially breaking and/or performance-sensitive changes.

@jakevdp jakevdp self-assigned this Dec 13, 2024
@jakevdp jakevdp marked this pull request as draft December 13, 2024 00:17
@jakevdp jakevdp force-pushed the xla-abstractify branch 2 times, most recently from fd24bf9 to 2b7dbc5 Compare December 17, 2024 00:04
@jakevdp jakevdp changed the title Cleanup: merge core.concrete_aval & xla.abstractify Cleanup: toward merging core.concrete_aval & xla.abstractify Dec 17, 2024
@jakevdp jakevdp marked this pull request as ready for review December 17, 2024 17:28
@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Dec 17, 2024
@copybara-service copybara-service bot merged commit 0fa5419 into jax-ml:main Dec 17, 2024
22 of 23 checks passed
@jakevdp jakevdp deleted the xla-abstractify branch December 17, 2024 20:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants