Move iter_children and iter_modules to functions

[samanklesaria]

Closes #4943

[vfdev-5]

Thanks!
I left few comments for improvements.
Apply the comments to both nnx.iter_modules and nnx.iter_children.
Please also expose new methods in the docs_nnx: https://github.com/google/flax/blob/main/docs_nnx/api_reference/flax.nnx/module.rst
and also in __init__ to be addressable as nnx.*.

However, we should confirm the submodule where to put these methods.
@cgarciae flax/nnx/module.py is ok or in graph.py ?

[cgarciae]

@vfdev-5 I think they should remain in module.py.

[vfdev-5]

LGTM, thanks @samanklesaria