[nnx] add pytree support

[cgarciae]

What does this PR do?

Add support for Pytree nodes in graph_utils

[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: Patch coverage is 93.93939% with 6 lines in your changes are missing coverage. Please review.

Project coverage is 58.70%. Comparing base (0c7d256) to head (c898774).
Report is 3 commits behind head on main.

Files	Patch %	Lines
flax/experimental/nnx/nnx/graph_utils.py	90.62%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3732      +/-   ##
==========================================
+ Coverage   58.43%   58.70%   +0.27%     
==========================================
  Files         102      102              
  Lines       12365    12412      +47     
==========================================
+ Hits         7225     7286      +61     
+ Misses       5140     5126      -14     

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