[nnx] non-str State keys

[cgarciae]

What does this PR do?

• Graph node can now emit non-string keys, this helps to sort int keys such as to from list and tuple. In general they can emit any type that implement the Key protocol, which is any Hashable object that also implements __lt__.
• State now contains Key keys.
• State.flat_state now returns a dict[tuple[Key, ...], Any] instead of dict[str, Any] with keys as / separated strings.
• State.from_flat_path now creates a State from a dict[tuple[Key, ...], Any].
• _graph_node_* functions are now methods of GraphNode.
• List now emits int keys. To implement the graph node protocol correctly it overrides a couple of _graph_node_* methods.
• Fixed graph_pop bug.

[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 94.26230% with 7 lines in your changes are missing coverage. Please review.

Project coverage is 60.34%. Comparing base (cc740d4) to head (1d5f718).

Files	Patch %	Lines
flax/experimental/nnx/nnx/graph_utils.py	94.44%	3 Missing ⚠️
flax/experimental/nnx/nnx/helpers.py	76.92%	3 Missing ⚠️
flax/typing.py	80.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3802   +/-   ##
=======================================
  Coverage   60.34%   60.34%           
=======================================
  Files         101      101           
  Lines       12862    12875   +13     
=======================================
+ Hits         7761     7769    +8     
- Misses       5101     5106    +5     

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