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.

Sign up for GitHub

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

Jump to bottom

Fix performance bug in vmap call #230

Merged
merged 2 commits into from
Feb 4, 2025
Merged

Fix performance bug in vmap call #230

merged 2 commits into from
Feb 4, 2025

Conversation

jfeser
Copy link
Contributor

@jfeser jfeser commented Feb 4, 2025

vmap calls repr on its function argument, which can be slow when we pass a term containing a large tensor. Avoid this cost with a wrapper that has a trivial repr.

@jfeser
Copy link
Contributor Author

jfeser commented Feb 4, 2025

With this change, our slowest test goes from ~12s to <1s.

@jfeser jfeser requested a review from eb8680 February 4, 2025 20:27
@jfeser jfeser merged commit a400a20 into master Feb 4, 2025
3 checks passed
@jfeser jfeser deleted the jf-vmap-perf branch February 4, 2025 20:51
@jfeser jfeser mentioned this pull request Feb 4, 2025
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@eb8680 eb8680 eb8680 approved these changes

Assignees
No one assigned
Labels
status:awaiting review
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@jfeser @eb8680