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

Tuple arity mismatch: %d != %d; tuple: %s.

Error message

Tuple arity mismatch: %d != %d; tuple: %s.

What it means

During PyTreeDef::FlattenUpTo in JAX's pytree C++ extension, the object at this position is an exact tuple but its length does not equal the arity recorded in the treedef node. JAX enforces strict structural equality of prefixes, so a tuple of the wrong length is rejected even if elements are compatible.

Source

Thrown at jaxlib/pytree.cc:1008

              "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(
              absl::StrFormat("Expected list, got %s.",
                              nb::cast<std::string_view>(nb::repr(object))));
        }
        nb::list list = nb::borrow<nb::list>(object);
        if (list.size() != node.arity) {
          throw std::invalid_argument(absl::StrFormat(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix the tuple length at that position to match the treedef arity shown in the message (%d != %d; first is the object's size).
2. If the data legitimately changed size, do not reuse the old treedef: re-trace/re-jit or recompute the treedef from the new object.
3. Build tuples deterministically (e.g. tuple(fixed_length_iterable)) instead of conditional appends.
4. For variable-length data, use a registered custom pytree node that stores length in aux_data or use lists/dicts with matching structure.

Example fix

# before
treedef = jax.tree_util.tree_structure(((0, 0), (0, 0)))
obj = ((1, 2), (3,))          # inner tuple arity 1 != 2
# after
obj = ((1, 2), (3, 4))        # arity matches

node_arity = 2  # from treedef / error message
assert len(my_tuple) == node_arity, f'expected {node_arity}, got {len(my_tuple)}'

def tuple_matches(t: tuple[int, ...], arity: int) -> bool:
    return type(t) is tuple and len(t) == arity

try:
    treedef.flatten_up_to(obj)
except ValueError as e:
    if 'Tuple arity mismatch' in str(e):
        raise ValueError(f'reshape input to arity in message: {e}') from e
    raise

When it happens

Trigger: Flattening an object against a prefix treedef whose tuple node has arity N while the runtime tuple has M != N elements; e.g. jax.jit argument donated/reused mismatch, tree_unflatten-then-modify, calling FlattenUpTo-based internal paths (jit cache hits, vmap in_axes/out_axes traversal) with a shortened/lengthened tuple.

Common situations: Appending/removing an element from a config tuple between runs while reusing a compiled/cached function, off-by-one tuple construction (trailing comma adding a 1-tuple), conditional code paths building tuples of different lengths, pickled treedefs from older data shape.

Related errors

• Effects not supported in partial-eval of `checkpoint`/`remat
• Expected tuple, got %s.
• List arity mismatch: %d != %d; list: %s.
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64

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