flax

90715bef - [nnx] disallow Array leaves

[nnx] disallow Array leaves Numpy and JAX Array's are no longer consider state leaves. This makes the structure of the State completely determined by Variables, which apart from being more predictable it produces structural stability invariant to leaf type changes which let to issues such as https://github.com/google/flax/issues/4142. ```python class Foo(nnx.Module): def __init__(self): self.a = jnp.array(1) # no longer allowed, instead... self.b = nnx.Param(jnp.array(1)) # just use Variables ``` Also migrates all remaining tests from pytest to absl to ensure they are tested correctly internally. PiperOrigin-RevId: 671372717