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

added wrt option to nnx.Optimizer #3983

Merged
merged 1 commit into from
Jun 12, 2024
Merged

Conversation

chiamp
Copy link
Collaborator

@chiamp chiamp commented Jun 11, 2024

added wrt option to nnx.Optimizer, to allow the user to specify which variables to keep track off in the optimizer state and which variables update in nnx.Optimizer.update.

@chiamp chiamp self-assigned this Jun 11, 2024
@chiamp chiamp marked this pull request as draft June 11, 2024 19:07
@chiamp chiamp changed the title added wrt option to nnx.Optimizer.update added wrt option to nnx.Optimizer.update Jun 11, 2024
@chiamp chiamp force-pushed the optimizer branch 2 times, most recently from 7951d51 to 7e43f9f Compare June 11, 2024 22:23
@codecov-commenter
Copy link

codecov-commenter commented Jun 11, 2024

Codecov Report

Attention: Patch coverage is 0% with 24 lines in your changes missing coverage. Please review.

Project coverage is 0.00%. Comparing base (31adb00) to head (9a40b3f).
Report is 7 commits behind head on main.

Files Patch % Lines
flax/nnx/tests/optimizer_test.py 0.00% 18 Missing ⚠️
flax/nnx/nnx/training/optimizer.py 0.00% 5 Missing ⚠️
flax/nnx/__init__.py 0.00% 1 Missing ⚠️
Additional details and impacted files
@@          Coverage Diff          @@
##            main   #3983   +/-   ##
=====================================
  Coverage   0.00%   0.00%           
=====================================
  Files        106     106           
  Lines      13582   13633   +51     
=====================================
- Misses     13582   13633   +51     

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

@chiamp chiamp marked this pull request as ready for review June 11, 2024 22:58
@chiamp chiamp changed the title added wrt option to nnx.Optimizer.update added wrt option to nnx.Optimizer Jun 11, 2024
@chiamp chiamp requested a review from cgarciae June 12, 2024 05:25
@copybara-service copybara-service bot merged commit 56216fa into google:main Jun 12, 2024
19 checks passed
@chiamp chiamp deleted the optimizer branch June 12, 2024 21:12
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