[nnx] Add support for python container types

[cgarciae]

What does this PR do?

• Adds support for list, tuple, and dict as Graph node types.
• Renames ModuleDef as GraphDef
• Refactors all the graph related code into graph_utils.py

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov-commenter]

Codecov Report

Attention: 47 lines in your changes are missing coverage. Please review.

Comparison is base (091df13) 52.10% compared to head (eef8066) 52.96%.

Files	Patch %	Lines
flax/experimental/nnx/nnx/graph_utils.py	88.16%	40 Missing ⚠️
flax/experimental/nnx/nnx/proxy_caller.py	82.85%	6 Missing ⚠️
flax/experimental/nnx/nnx/helpers.py	80.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3486      +/-   ##
==========================================
+ Coverage   52.10%   52.96%   +0.86%     
==========================================
  Files          91       94       +3     
  Lines       10968    11199     +231     
==========================================
+ Hits         5715     5932     +217     
- Misses       5253     5267      +14     

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