[nnx] add submodule iterator #3581
Conversation
Codecov ReportAttention:
Additional details and impacted files@@ Coverage Diff @@
## main #3581 +/- ##
==========================================
+ Coverage 56.31% 56.46% +0.15%
==========================================
Files 100 100
Lines 11973 11994 +21
==========================================
+ Hits 6742 6773 +31
+ Misses 5231 5221 -10 ☔ View full report in Codecov by Sentry. |
| @@ -482,29 +482,10 @@ def sow( | |||
| reduced_value = reduce_fn(init_fn(), value) | |||
| setattr(self, name, variable_type(reduced_value)) | |||
|
|
|||
| def for_each( | |||
There was a problem hiding this comment.
Should this be called submodules and not module? Or better yet iter_submodules?
There was a problem hiding this comment.
Called it modules() as it would be familiar to Pytorch users. Wondering if we should try to follow their conventions when possible?
| def iter_nodes(node: tp.Any) -> tp.Iterator[tuple[Path, tp.Any]]: | ||
| visited: set[int] = set() | ||
| path_parts: PathParts = () | ||
| yield from _iter_nodes(node, visited, path_parts) | ||
|
|
||
|
|
||
| def _iter_nodes( | ||
| node: tp.Any, visited: set[int], path_parts: PathParts |
There was a problem hiding this comment.
can the node args have a more specific type annotation than tp.Any?
There was a problem hiding this comment.
problem is that a node could be of any type that is registered.
cgarciae
force-pushed
the
nnx-submodule-iterator
branch
from
43eb65a
to
b220ec4
Compare
cgarciae
force-pushed
the
nnx-submodule-iterator
branch
from
b220ec4
to
1e75509
Compare
|
@superbobry added test. |
What does this PR do?
Adds method to iterator over all unique submodules.