Experimental-pytree flag causes crash

[NeilGirdhar]

from dataclasses import dataclass, field
from typing import Any, dataclass_transform, override

import jax.numpy as jnp
from flax import nnx
from jax import Array, vmap


@dataclass_transform(field_specifiers=(field,))
class DataClassModule(nnx.Module):
    @override
    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs, experimental_pytree=True)
        dataclass()(cls)


class SomeModule(nnx.Module):
    def __init__(self, epsilon: Array):
        super().__init__()
        self.epsilon = epsilon


class SomeDataclassModule(DataClassModule):
    def __init__(self) -> None:
        super().__init__()
        self.sm = SomeModule(jnp.zeros(1))


def f(m: SomeDataclassModule, x: Array) -> None:
    pass


module = SomeDataclassModule()
z = jnp.zeros(10)
vmap(f, in_axes=(None, 0))(module, z)

gives

Traceback (most recent call last):
  File "/home/neil/src/cmm/a.py", line 35, in <module>
    vmap(f, in_axes=(None, 0))(module, z)
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/jax/_src/api.py", line 1221, in vmap_f
    in_axes_flat = flatten_axes("vmap in_axes", in_tree, (in_axes, 0), kws=True)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/jax/_src/api_util.py", line 400, in flatten_axes
    dummy = tree_unflatten(treedef, [object()] * treedef.num_leaves)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/jax/_src/tree_util.py", line 84, in tree_unflatten
    return treedef.unflatten(leaves)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/flax/nnx/nnx/module.py", line 444, in _module_unflatten
    return graph.merge(graphdef, State(zip(paths, variables)))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/flax/nnx/nnx/graph.py", line 1306, in merge
    node, _ = unflatten(graphdef, state)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/flax/nnx/nnx/graph.py", line 471, in unflatten
    node = _graph_unflatten(
           ^^^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/flax/nnx/nnx/graph.py", line 628, in _graph_unflatten
    children = _get_children()
               ^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/flax/nnx/nnx/graph.py", line 569, in _get_children
    children[key] = _graph_unflatten(
                    ^^^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/flax/nnx/nnx/graph.py", line 628, in _graph_unflatten
    children = _get_children()
               ^^^^^^^^^^^^^^^
  File "/home/neil/src/cmm/.venv/lib/python3.12/site-packages/flax/nnx/nnx/graph.py", line 575, in _get_children
    raise ValueError(f'Expected a leaf for {key!r}, but got {value!r}')
ValueError: Expected a leaf for 'epsilon', but got <object object at 0x7ccaa815f050>

Assigning a float to epsilon makes the problem disappear.

(Tested on main and latest.)

[cgarciae]

Interesting counter examples for Module's pytree definition. The issue is that we decide what is static based on the type.

[NeilGirdhar]

Okay, I understand. Feel free to close if you like.