maintain data/static definition in split / state

[copybara-service]

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

    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