search

google / flax

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

Added `nnx.graph.iter_children` #3991

Closed chiamp closed 3 months ago

chiamp commented 3 months ago

Added nnx.graph.iter_children, which returns an iterator of only immediate child Module's (contrasted with nnx.graph.iter_modules that returns an iterator of ALL nested Module's). Also updated docstring and tests for both methods. Maybe we can use more informative names like iter_all_modules and iter_child_modules?

codecov-commenter commented 3 months ago

Codecov Report

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

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

Files Patch % Lines
flax/nnx/tests/module_test.py 0.00% 26 Missing :warning:
flax/nnx/nnx/module.py 0.00% 6 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## main #3991 +/- ## ===================================== Coverage 0.00% 0.00% ===================================== Files 106 106 Lines 13582 13657 +75 ===================================== - Misses 13582 13657 +75 ```

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