[nnx] switch to nested State representation

[cgarciae]

What does this PR do?

Uses a nested representation instead of a flattened path representation.

State({
  'linear': {
    'kernel': Array([[0.31696808, 0.55285215],
           [0.31418085, 0.7399571 ]], dtype=float32),
    'bias': Array([0., 0.], dtype=float32)
  }
})

The PR also implements State.__getattr__ and State.__setattr__ to allow for a more natural access mechanism. In particular, since paths are preserved exactly from the corresponding Module, you get the exact same access syntax.

# getattr
model.linear.kernel == state.linear.kernel
# getitem
model.linear.kernel == state['linear']['kernel']

Integer access is also allowed keep the 1-1 access correspondence with list and tuple nodes:

class Foo(nnx.Module):
  def __init__(self, *, rngs: nnx.Rngs):
    self.layers = [
      nnx.Linear(1, 2, rngs=rngs),
      nnx.Linear(2, 3, rngs=rngs),
    ]

model = Foo(rngs=nnx.Rngs(0))
state = model.get_state()

assert model.layers[0].kernel.shape == (1, 2)
assert state.layers[0].kernel.shape == (1, 2)
assert model.layers[1].kernel.shape == (2, 3)
assert state.layers[1].kernel.shape == (2, 3)

[codecov-commenter]

Codecov Report

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

Comparison is base (70214f4) 52.98% compared to head (096adc8) 53.28%.

Files	Patch %	Lines
flax/experimental/nnx/nnx/state.py	76.27%	14 Missing ⚠️
flax/experimental/nnx/nnx/graph_utils.py	92.85%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3502      +/-   ##
==========================================
+ Coverage   52.98%   53.28%   +0.30%     
==========================================
  Files          94       95       +1     
  Lines       11206    11238      +32     
==========================================
+ Hits         5937     5988      +51     
+ Misses       5269     5250      -19     

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

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB