jax-ml/jax · error · std::invalid_argument

Expected named tuple, got %s.

Error message

Expected named tuple, got %s.

What it means

JAX's FlattenUpTo reached a kNamedTuple node, but the object at that position is not a tuple subclass with a _fields attribute (i.e. not a collections.namedtuple/typing.NamedTuple). Plain tuples, lists, and dataclasses fail this check.

Source

Thrown at jaxlib/pytree.cc:1062

        if (!IsSortedPyDictKeysEqual(keys, node.sorted_dict_keys)) {
          // Convert to a nb::list for nb::repr to avoid having to stringify a
          // vector. This is error path so it is fine to pay conversion cost.
          throw std::invalid_argument(absl::StrFormat(
              "Dict key mismatch; expected keys: %s; present keys: %s.",
              nb::cast<std::string_view>(
                  nb::repr(nb::cast(node.sorted_dict_keys))),
              nb::cast<std::string_view>(nb::repr(nb::cast(keys)))));
        }
        for (nb::handle key : keys) {
          agenda.push_back(dict[key]);
        }
        break;
      }

      case PyTreeKind::kNamedTuple: {
        if (!nb::isinstance<nb::tuple>(object) ||
            !nb::hasattr(object, "_fields")) {
          throw std::invalid_argument(
              absl::StrFormat("Expected named tuple, got %s.",
                              nb::cast<std::string_view>(nb::repr(object))));
        }
        nb::tuple tuple = nb::borrow<nb::tuple>(object);
        if (tuple.size() != node.arity) {
          throw std::invalid_argument(absl::StrFormat(
              "Named tuple arity mismatch: %d != %d; tuple: %s.", tuple.size(),
              node.arity, nb::cast<std::string_view>(nb::repr(object))));
        }
        if (tuple.type().not_equal(node.node_data)) {
          throw std::invalid_argument(absl::StrFormat(
              "Named tuple type mismatch: expected type: %s, tuple: %s.",
              nb::cast<std::string_view>(nb::repr(node.node_data)),
              nb::cast<std::string_view>(nb::repr(object))));
        }
        for (nb::handle entry : tuple) {
          agenda.push_back(nb::borrow<nb::object>(entry));
        }

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass an instance of the same namedtuple type at that position.
2. Re-derive the treedef after switching container kinds (dataclass vs namedtuple): register dataclasses with jax.tree_util.register_dataclass or use flax.struct.dataclass.
3. Keep a single canonical type definition module to avoid duplicate namedtuple definitions with identical names but different identity.
4. If you need struct-like flexibility, use pytree-registered dataclasses instead of namedtuples.

Example fix

# before
from collections import namedtuple
P = namedtuple('P', ['x', 'y'])
treedef = jax.tree_util.tree_structure(P(1, 2))
obj = (3, 4)                       # plain tuple -> error
# after
obj = P(3, 4)                      # same namedtuple type

from collections import namedtuple
assert hasattr(obj, '_fields') and isinstance(obj, tuple), f'expected namedtuple, got {type(obj).__name__}'

import typing
def is_namedtuple_instance(x) -> bool:
    return isinstance(x, tuple) and hasattr(x, '_fields')

try:
    treedef.flatten_up_to(obj)
except (ValueError, TypeError) as e:
    if 'Expected named tuple' in str(e):
        obj = MyNamedTuple(*obj)  # rebuild as the traced namedtuple type
    else:
        raise

When it happens

Trigger: Flattening against a treedef that captured a namedtuple while the runtime value is a plain tuple, list, or dataclass instance; occurs when converting namedtuples to dataclasses or plain tuples during refactors while reusing treedefs/compiled functions, or when a namedtuple import path resolves to a different object.

Common situations: Migrating NamedTuple configs to @dataclass without re-tracing, replacing namedtuple(a, b)(...) with (a, b) for brevity, namedtuple field renaming creating new types, pickling treedefs across code versions.

Related errors

• Python int {value} too large to convert to int32
• The names should be exclusive and should not intersect in `n
• Expected tuple, got %s.
• Expected list, got %s.
• Expected dict, got %s.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/1ba987a791e4a186. Report an issue: GitHub.