[NNX] Make `State`'s mapping by reference to allow attribute add/delete

[IvyZX]

This allows adding and deleting entries from the nnx.State mappings, which is crucial for flexible model surgery.

[codecov-commenter]

Codecov Report

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

Project coverage is 0.00%. Comparing base (2c7d7cd) to head (a8e8045). Report is 44 commits behind head on main.

Files	Patch %	Lines
flax/experimental/nnx/tests/test_state.py	0.00%	12 Missing :warning:
flax/experimental/nnx/nnx/state.py	0.00%	6 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##             main   #3930       +/-   ##
==========================================
- Coverage   60.43%   0.00%   -60.44%     
==========================================
  Files         105     101        -4     
  Lines       13263   13039      -224     
==========================================
- Hits         8015       0     -8015     
- Misses       5248   13039     +7791     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.