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

jnp.einsum: default to optimize='auto' #25214

Merged
merged 1 commit into from
Jan 6, 2025

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Dec 2, 2024

Fixes #24929

This will lead to much faster tracing of jnp.einsum as the number of arguments grow. This may lead to less efficient runtime performance in some cases. I think avoiding unconditional exponential trace-time is worth the potential runtime tradeoff in the case of many arguments.

@jakevdp jakevdp added the pull ready Ready for copybara import and testing label Dec 2, 2024
@jakevdp jakevdp self-assigned this Dec 2, 2024
@jakevdp jakevdp mentioned this pull request Dec 2, 2024
@jakevdp jakevdp requested a review from hawkinsp January 6, 2025 16:11
Copy link
Collaborator

@hawkinsp hawkinsp left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd mention this in the changelog, but LGTM.

@jakevdp
Copy link
Collaborator Author

jakevdp commented Jan 6, 2025

I'd mention this in the changelog, but LGTM.

Done, thanks!

@copybara-service copybara-service bot merged commit 52cc5c7 into jax-ml:main Jan 6, 2025
20 of 23 checks passed
@jakevdp jakevdp deleted the einsum-optimize branch January 6, 2025 19:38
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.

Small Einsum is hanging
3 participants