Add NNX

[cgarciae]

What does this PR do?

Adds cgarciae/nnx as an experimental submodule from Flax:

from flax.experimental import nnx

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov-commenter]

Codecov Report

Merging #3218 (8e9df27) into main (cbf7bea) will decrease coverage by 31.46%.
Report is 10 commits behind head on main.
The diff coverage is 86.31%.

@@             Coverage Diff             @@
##             main    #3218       +/-   ##
===========================================
- Coverage   83.62%   52.16%   -31.46%     
===========================================
  Files          56       91       +35     
  Lines        6790    10954     +4164     
===========================================
+ Hits         5678     5714       +36     
- Misses       1112     5240     +4128     

Files	Coverage Δ
flax/experimental/nnx/__init__.py	100.00% <100.00%> (ø)
flax/experimental/nnx/nnx/errors.py	100.00% <100.00%> (ø)
flax/experimental/nnx/nnx/nn/activations.py	100.00% <100.00%> (ø)
flax/experimental/nnx/tests/test_compatibility.py	100.00% <100.00%> (ø)
flax/experimental/nnx/tests/test_containers.py	100.00% <100.00%> (ø)
flax/experimental/nnx/tests/test_helpers.py	100.00% <100.00%> (ø)
flax/experimental/nnx/tests/test_ids.py	100.00% <100.00%> (ø)
flax/experimental/nnx/tests/test_integration.py	100.00% <100.00%> (ø)
flax/experimental/nnx/tests/test_partitioning.py	100.00% <100.00%> (ø)
flax/experimental/nnx/tests/test_rngs.py	100.00% <100.00%> (ø)
... and 27 more

... and 41 files with indirect coverage changes

[8bitmp3]

🚀 @cgarciae @levskaya