[nnx] improved graph update mechanism

[cgarciae]

What does this PR do?

Adds a mechanism that allows performing arbitrary mutations on a cloned graph and recreating them on the original graph, this is useful to faithfully update a graph across jax boundaries.

Changes

• adds MutableNodeImpl.clear API to enable graph_unflatten to update existing references
• graph_flatten now returns a ref_to_index: Mapping[Any, int] that maps the node references to their flatten index.
• graph_unflatten now returns a index_to_ref: dict[int, Any] maps index flatten references to their node reference.
• graph_unflatten now accepts a ref_cache: dict[int, Any] that maps flatten indexes to existing node references to use instead of creating new ones. clear is used to empty the existing node's state so they can be treated as new nodes during the unflattening process.

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 89.85507% with 21 lines in your changes are missing coverage. Please review.

Project coverage is 59.49%. Comparing base (83d118a) to head (10a6b2d).
Report is 7 commits behind head on main.

Files	Patch %	Lines
flax/experimental/nnx/nnx/graph_utils.py	84.95%	17 Missing ⚠️
flax/experimental/nnx/nnx/variables.py	33.33%	4 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3759      +/-   ##
==========================================
+ Coverage   58.75%   59.49%   +0.74%     
==========================================
  Files         101      101              
  Lines       12413    12621     +208     
==========================================
+ Hits         7293     7509     +216     
+ Misses       5120     5112       -8     

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

[superbobry]

Nit: I'd expect the type checker to be able to infer this type with no hints.

If you decide to remove this declaration, please remove other similar ones in this file as well.

[cgarciae]

This is here to force the static checker to treat static as a GraphDef everywhere, else its reset to Any when its redefined:

static, idx_out_idx_in = static_out.value

[superbobry]

Nit: I think it might be more clear to do

state, static, _ = ...

instead, since you already do unpacking.