moved `Module.iter_*` methods to `nnx.graph`

google / flax

Flax is a neural network library for JAX that is designed for flexibility.

https://flax.readthedocs.io

Apache License 2.0

5.78k stars 610 forks source link

Open chiamp opened 2 weeks ago

chiamp commented 2 weeks ago

Moved Module.iter_* methods to nnx.graph (see nnx.graph.iter_nodes and nnx.graph.iter_child_nodes for more detail).

~Also added a filter warning for a deprecation error. See more detail here.~