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 nnx.graph.iter_children #3991

Merged
merged 1 commit into from
Jun 13, 2024
Merged

Added nnx.graph.iter_children #3991

merged 1 commit into from
Jun 13, 2024

Conversation

chiamp
Copy link
Collaborator

@chiamp chiamp commented Jun 12, 2024

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?

@chiamp chiamp self-assigned this Jun 12, 2024
@chiamp chiamp changed the title added nnx.graph.iter_children added nnx.graph.iter_children Jun 12, 2024
@chiamp chiamp changed the title added nnx.graph.iter_children dded nnx.graph.iter_children Jun 12, 2024
@chiamp chiamp changed the title dded nnx.graph.iter_children Added nnx.graph.iter_children Jun 12, 2024
@codecov-commenter
Copy link

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 ⚠️
flax/nnx/nnx/module.py 0.00% 6 Missing ⚠️
Additional details and impacted files
@@          Coverage Diff          @@
##            main   #3991   +/-   ##
=====================================
  Coverage   0.00%   0.00%           
=====================================
  Files        106     106           
  Lines      13582   13657   +75     
=====================================
- Misses     13582   13657   +75     

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

linear Linear
submodule SubModule
"""
node_dict = graph.get_node_impl(self).node_dict(self)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice! This is more general.

@copybara-service copybara-service bot merged commit e6b83ad into google:main Jun 13, 2024
19 checks passed
@chiamp chiamp deleted the iter branch June 13, 2024 22:43
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