flax

404a4208 - maintain data/static definition in split / state

maintain data/static definition in split / state `nnx.graph` APIs now follow the data / static definitions for `nnx.Pytree` instances. Previously `nnx.graph.flatten` match the default definitions for Pytree most of the time but deviated on edcases, now they always match. Example ```python class Foo(nnx.Pytree): def __init__(self, data, static): self.data = nnx.data(data) self.static = nnx.static(static) tree = Foo(1, 2) state = nnx.state(tree) # previously this was false assert 'data' in state ``` PiperOrigin-RevId: 868239931