-
Notifications
You must be signed in to change notification settings - Fork 601
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
Make FlatState a Mapping instead of a dict #3928
Conversation
@cgarciae (Done, as you requested.) |
fa0436e
to
6776c47
Compare
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #3928 +/- ##
==========================================
- Coverage 60.43% 0.00% -60.44%
==========================================
Files 105 105
Lines 13263 13328 +65
==========================================
- Hits 8015 0 -8015
- Misses 5248 13328 +8080 ☔ View full report in Codecov by Sentry. |
@NeilGirdhar looks great! Thanks for doing this. |
Seems there are some git conflicts. |
Rebased to main. |
* Add nnx.traversals.{flatten_mapping, unflatten_dict} * These are modified from traverse_util.{flatten_dict, unflatten_dict} * They are annotated so that any future changes to the function (e.g., changing it back to work with dicts only) triggers a type error. * Minor tweaks to imports: * import from collections.abc instead of typing since the latter imports are deprecated. * Import from flax.typing instead of flax.core.scope since the latter has been moved. * Add from __future__ import annotations so that the annotations work on Python 3.9. * Minor tweaks to code: * Annotate some private functions to make them easier to understand. * When printing a type, print its __qualname__ since that's a bit easier to read (str instead of <class 'str'>). Fixes google#3879
Add nnx.traversals.{flatten_mapping, unflatten_dict}
Minor tweaks to imports:
__future__
import annotations so that the annotations work on Python 3.9.Minor tweaks to code:
Fixes #3879