search

google / flax

Flax is a neural network library for JAX that is designed for flexibility.
https://flax.readthedocs.io
Apache License 2.0
5.79k stars 609 forks source link

added `wrt` option to `nnx.Optimizer` #3983

Closed chiamp closed 3 weeks ago

chiamp commented 3 weeks ago

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.

codecov-commenter commented 3 weeks ago

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 :warning:
flax/nnx/nnx/training/optimizer.py 0.00% 5 Missing :warning:
flax/nnx/__init__.py 0.00% 1 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #3983 +/- ## ===================================== Coverage 0.00% 0.00% ===================================== Files 106 106 Lines 13582 13633 +51 ===================================== - Misses 13582 13633 +51 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.