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

Expected tuple, got %s.

Error message

Expected tuple, got %s.

What it means

JAX/pytree raised this while flattening a Python object against a stored PyTreeDef (treedef) via FlattenUpTo. The treedef records that this node in the tree structure is an exact built-in tuple, but the object being flattened has a different type at that position. It is thrown whenever a function/transform re-applies a structure captured from different data (e.g. jit/pmap with donated or reused arguments, tree_unflatten-like paths, vmap out_axes matching).

Source

Thrown at jaxlib/pytree.cc:1002

        --leaf;
        break;

      case PyTreeKind::kNone:
        if (!object.is_none()) {
          throw std::invalid_argument(absl::StrFormat(
              "Expected None, got %s.\n\n"
              "In previous releases of JAX, flatten-up-to used to "
              "consider None to be a tree-prefix of non-None values. To obtain "
              "the previous behavior, you can usually write:\n"
              "  jax.tree.map(lambda x, y: None if x is None else f(x, y), a, "
              "b, is_leaf=lambda x: x is None)",
              nb::cast<std::string_view>(nb::repr(object))));
        }
        break;

      case PyTreeKind::kTuple: {
        if (!PyTuple_CheckExact(object.ptr())) {
          throw std::invalid_argument(
              absl::StrFormat("Expected 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(
              "Tuple arity mismatch: %d != %d; tuple: %s.", tuple.size(),
              node.arity, nb::cast<std::string_view>(nb::repr(object))));
        }
        for (nb::handle entry : tuple) {
          agenda.push_back(nb::borrow<nb::object>(entry));
        }
        break;
      }

      case PyTreeKind::kList: {
        if (!PyList_CheckExact(object.ptr())) {
          throw std::invalid_argument(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the object an exact built-in tuple at that position: convert with tuple(x) before passing it in.
2. Re-derive the treedef from the current object (re-jit / re-flatten) instead of reusing a treedef captured from earlier data.
3. If you intentionally changed structure, clear cached compilations (e.g. new jax.jit wrapper) or avoid reusing stored treedefs across structure changes.
4. Replace tuple nodes with a structure-insensitive container (e.g. dict or custom pytree node registered with jax.tree_util.register_pytree_node) if you need flexible sequence types.

Example fix

# before
leaves = treedef.flatten_up_to((1, [2, 3]))  # second element was tuple in treedef
# TypeError/ValueError: Expected tuple, got [2, 3].

# after
leaves = treedef.flatten_up_to((1, (2, 3)))  # use exact tuple

import jax
expected = jax.tree_util.tree_structure(reference_tuple)
def structure_ok(obj):
    try:
        jax.tree_util.tree_flatten(obj)
        return True
    except (TypeError, ValueError):
        return False

from jax.tree_util import tree_structure
def is_exact_tuple_tree(obj, ref):
    return tree_structure(obj) == tree_structure(ref) and all(
        isinstance(x, tuple) and type(x) is tuple
        for x in jax.tree_util.tree_leaves(obj) or [obj]
    )

try:
    treedef.flatten_up_to(obj)
except (ValueError, TypeError) as e:
    if 'Expected tuple' in str(e):
        obj = jax.tree_util.tree_map(lambda x: tuple(x) if isinstance(x, list) else x, obj)
    else:
        raise

When it happens

Trigger: Calling an API that flattens a tree up to a prefix treedef (e.g. jax.tree_util.tree_flatten(..., is_leaf=...), FlattenUpTo paths used by jax.jit cache matching / tree_map with mismatched structures) where the treedef node is kTuple (PyTuple_CheckExact) but the runtime object is a list, namedtuple, or other sequence. Subclassing tuple is not enough: the check is exact.

Common situations: Passing a list where a tuple was captured at trace time (jit re-trace vs cached path), refactoring data containers from tuple to list (or dataclass) between code revisions, mixing namedtuple and plain tuple, replaying a treedef saved/pickled from an older run against new data.

Related errors

• Effects not supported in partial-eval of `checkpoint`/`remat
• Tuple arity mismatch: %d != %d; tuple: %s.
• Expected list, got %s.
• Expected dict, got %s.
• Expected named tuple, got %s.

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