Skip to content

Optimize struct.dataclass #3024

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

Merged
merged 1 commit into from
Apr 20, 2023
Merged

Optimize struct.dataclass #3024

merged 1 commit into from
Apr 20, 2023

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Apr 11, 2023

What does this PR do?

Currently struct.dataclass's implementation of pytree unflattening uses the class constructor, this can be up to 10x slower than other libraries as the number of leaves grows:

Screenshot 2023-04-19 at 3 49 23 PM

This change replaces the use of the class constructor with object.__new__, with this change struct.dataclass is faster than the other frameworks.

@cgarciae cgarciae force-pushed the optimize-pytreenode branch from 48e2ec3 to ab9051b Compare April 19, 2023 19:05
@cgarciae cgarciae changed the title Optimize PyTreeNode Optimize struct.dataclass Apr 19, 2023
@copybara-service copybara-service bot merged commit 7c9b3b3 into main Apr 20, 2023
@copybara-service copybara-service bot deleted the optimize-pytreenode branch April 20, 2023 12:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants