nnx.data is lost during nnx.eval_shape

[qGentry]

Hey folks,

I've just noticed that for some reason data registered via nnx.data in the module is lost during nnx.eval_shape. I'm currently in the process of migrating our codebase to Flax.NNX and instead of rewriting entire metrics implementation in NNX style I just bundle metrics state in NNX module as a nnx.data field. The issue I've faced is that for some reason if I construct abstract shape with nnx.eval_shape, then it is lost for some reason.

Here's a tiny repro:

from flax import nnx
import jax.numpy as jnp


class Module(nnx.Module):
    def __init__(self):
        self.data = nnx.data({"a": jnp.ones((8, 8))})


def _init_module() -> Module:
    module = Module()
    return module

module = _init_module()
print("State with direct init:")
print(nnx.state(module))
module = nnx.eval_shape(_init_module)
print("State with nnx.eval_shape:")
print(nnx.state(module))

Output:

State with direct init:
State({
  'data': {
    'a': Array([[1., 1., 1., 1., 1., 1., 1., 1.],
           [1., 1., 1., 1., 1., 1., 1., 1.],
           [1., 1., 1., 1., 1., 1., 1., 1.],
           [1., 1., 1., 1., 1., 1., 1., 1.],
           [1., 1., 1., 1., 1., 1., 1., 1.],
           [1., 1., 1., 1., 1., 1., 1., 1.],
           [1., 1., 1., 1., 1., 1., 1., 1.],
           [1., 1., 1., 1., 1., 1., 1., 1.]], dtype=float32)
  }
})
State with nnx.eval_shape:
State({})

[qGentry]

I genuinely don't understand the logic behind the state and the nnx.data interaction. For example, I can construct proper state via plain jax.eval_shape:

from flax import nnx
import jax
import jax.numpy as jnp


class Module(nnx.Module):
    def __init__(self):
        self.data = nnx.data({"a": jnp.ones((8, 8))})


def _init_module() -> Module:
    return Module()


def _init_module_state() -> nnx.State:
    module = _init_module()
    return nnx.state(module)

gdef = nnx.graphdef(_init_module())

jax_state = jax.eval_shape(_init_module_state)
print("State with jax.eval_shape:")
print(jax_state)

module = nnx.merge(gdef, jax_state)
print("State after merging with graphdef:")
print(nnx.state(module))

But after the merge with the gdef, it is no longer included in nnx.state:

State with jax.eval_shape:
State({
  'data': {
    'a': ShapeDtypeStruct(shape=(8, 8), dtype=float32)
  }
})
State after merging with graphdef:
State({})

[vfdev-5]

Hi @qGentry , nnx.state works on nnx.Variables, arrays etc, but once created a module using nnx.eval_shape, its variables, arrays are replaced by ShapeDtypeStruct.

Split a graph node into a GraphDef and one or more States. State is a ``Mapping from strings or integers to Variables, Arrays or nested States.

https://flax.readthedocs.io/en/stable/api_reference/flax.nnx/graph.html#flax.nnx.split

[qGentry]

Hey @vfdev-5 , thanks for the clarification but I still feel like this is not consistent. Shouldn't in this case nnx.state work on ShapeDtypeStructs too? Otherwise it introduces those weird discrepancies like this.

[cgarciae]

I think we could make this work now that we have nnx.data, we need a mechanism that _graph_node_flatten can use to communicate to _graph_flatten that the element is explicitly a leaf or metadata.

flax/flax/nnx/pytreelib.py

Line 919 in 2204c82

def _graph_node_flatten(self):

Maybe _graph_node_flatten could just wrap elements in some private _Leaf(value) and _Metadata(value) types to let _graph_flatten know.

[samanklesaria]

Although I'm not 100% sure I'm understanding things correctly, here's how it seems things work.
Currently, there are two parallel flattening systems used within the flax project. There's the one used by Jax (with instances registered by jax.tree_util.register_pytree_with_keys) and there's the one used by Flax (which handles circular references and shared structure and other stuff that Jax can't deal with). But the Flax one isn't just an expanded version of the Jax one: the flax and Jax systems flatten to leaves in fundamentally incompatible ways.

Flattening with Jax
With jax, flattening follows a simple recursive pattern. To flatten a pytree:

• Get the node's children.
• Recursively flatten those children that are themselves pytrees. Leave non-pytree children alone.
• The result of this recursive process are the leaves.

Flax's Modules are registered as Pytrees with the Jax flattening system. In this system, the module's children are the parameters that are either A) not pytrees (like arrays or ints or ShapeDtypeStructs) or B) annotated with nnx.data.

This is the only use of nnx.data: for the jax flattening system.

Flattening with Flax
This is what is used when you call nnx.state. Flax's graph flattening system doesn't look at nnx.data declarations. Instead, the leaf nodes it finds when looking at a Module are just its descendent Variable and Array nodes. That's it: no numbers, no ShapeDtypeStructs, no nada. Adding nnx.data annotations aren't going to help you here: that's just for the Jax flattening system.

Assigning Dicts to Flax Module Attributes
If you want to assign a dictionary to an attribute of a flax Module like in your example and have it be treated as data by the flax flattening systems, this means you can't use nnx.data. Instead, you could use nnx.Variable.

from flax import nnx
import jax.numpy as jnp

class Module(nnx.Module):
    def __init__(self):
        self.data = nnx.Variable({"a": jnp.ones((8, 8))})

module = Module()
print("State with direct init:")
print(nnx.state(module))
module2 = nnx.eval_shape(Module)
print("State with nnx.eval_shape:")
print(nnx.state(module2))

Note that this is different from what the Pytree guide says https://flax.readthedocs.io/en/latest/guides/pytree.html. I think the guide oversimplifies the true ugliness of the situation.

[qGentry]

Instead, the leaf nodes it finds when looking at a Module are just its descendent Variable and Array nodes. That's it: no numbers, no ShapeDtypeStructs, no nada.

Is there any downsides to extend the list of values that is being considered with the ShapeDtypeStruct apart from only nnx.Variable & jax.Array? This generally seems like a reasonable idea because in JAX I'm always certain that jax.Arrays and jax.ShapeDtypeStruct can be used interchangeably in many contexts.