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

[nnx] TrainState uses struct #3788

Merged
merged 1 commit into from
Apr 1, 2024
Merged

Conversation

cgarciae
Copy link
Collaborator

@cgarciae cgarciae commented Mar 26, 2024

What does this PR do?

Reimplements nnx.TrainState using struct.PyTreeNode.

@cgarciae cgarciae force-pushed the nnx-train-state-struct branch 2 times, most recently from 8f4df26 to ee1a138 Compare March 26, 2024 18:04
Copy link

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@codecov-commenter
Copy link

codecov-commenter commented Mar 26, 2024

Codecov Report

Attention: Patch coverage is 93.54839% with 2 lines in your changes are missing coverage. Please review.

Project coverage is 60.29%. Comparing base (514c111) to head (51b97ed).

Files Patch % Lines
flax/experimental/nnx/nnx/helpers.py 91.30% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main    #3788      +/-   ##
==========================================
+ Coverage   60.16%   60.29%   +0.12%     
==========================================
  Files         101      101              
  Lines       12842    12846       +4     
==========================================
+ Hits         7727     7745      +18     
+ Misses       5115     5101      -14     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

@cgarciae cgarciae force-pushed the nnx-train-state-struct branch 2 times, most recently from 8014194 to 7b70786 Compare March 28, 2024 17:48
@cgarciae cgarciae force-pushed the nnx-train-state-struct branch 2 times, most recently from e4e436e to 7c3704e Compare March 29, 2024 12:13
@cgarciae cgarciae force-pushed the nnx-train-state-struct branch from 7c3704e to 51b97ed Compare March 30, 2024 23:31
@copybara-service copybara-service bot merged commit 9c4c8f5 into main Apr 1, 2024
21 checks passed
@copybara-service copybara-service bot deleted the nnx-train-state-struct branch April 1, 2024 13:06
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.

3 participants