[nnx] mutable array p2

[cgarciae]

What does this PR do?

• Adds nnx.mutable_array and nnx.MutableArray, re-exporting them from jax.experimental.
• Uses [...] notation in some places instead of .value.
• Adds __data__ to most / all public Object-derived types.
• __data__'s type is now tuple[str, ...] | Literal['all', 'auto']

[ravwojdyla]

👋 @cgarciae when I tried the 1st toy example (Functional API) using latest release 0.10.6 (directed from https://flax.readthedocs.io/en/latest/examples/core_examples.html#toy-examples). It failed with:

TypeError: JAX arrays are immutable and do not support in-place item assignment. Instead of x[idx] = y, use x = x.at[idx].set(y) or another .at[] method: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.ndarray.at.html

On line:

flax/examples/nnx_toy_examples/01_functional_api.py

Line 53 in 893a660

self.count[...] += 1

That line was changed in this PR. I noticed the same operation in the 2nd and 3rd example are ATM unchanged:

flax/examples/nnx_toy_examples/02_lifted_transforms.py

Line 54 in 893a660

self.count.value += 1

I can see that you have the part 3 PR in #4755, which is not updating the 2nd and 3rd examples.

Some questions for you:

• the failure on the latest release is as expected?
• should the examples 2 and 3 also be updated?

Btw, pending your response/direction, I'm happy to help 🙏

[Tomas542]

@ravwojdyla, hi. Did you installed flax from PyPI or build it from source? I think, Mutable Arrays changes not on the PyPI builds for now