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

List arity mismatch: %d != %d; list: %s.

Error message

List arity mismatch: %d != %d; list: %s.

What it means

PyTreeDef::FlattenUpTo in JAX detected a list node whose length differs from the arity stored in the treedef. Exact-length matching is required because the treedef's prefix must structurally match the object being flattened.

Source

Thrown at jaxlib/pytree.cc:1026

          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(
              "List arity mismatch: %d != %d; list: %s.", list.size(),
              node.arity, nb::cast<std::string_view>(nb::repr(object))));
        }
        for (nb::handle entry : list) {
          agenda.push_back(nb::borrow<nb::object>(entry));
        }
        break;
      }

      case PyTreeKind::kDict: {
        if (!PyDict_CheckExact(object.ptr())) {
          throw std::invalid_argument(
              absl::StrFormat("Expected dict, got %s.",
                              nb::cast<std::string_view>(nb::repr(object))));
        }
        nb::dict dict = nb::borrow<nb::dict>(object);
        std::vector<nb::object> keys = GetSortedPyDictKeys(dict.ptr());
        if (!IsSortedPyDictKeysEqual(keys, node.sorted_dict_keys)) {

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad/truncate the list to the arity reported in the message (object size vs node arity).
2. Wrap the jitted function so lists of varying length trigger re-trace appropriately (e.g. pass length as a jnp scalar arg, or convert to jnp arrays with batch dim) instead of relying on cached treedefs.
3. If lengths vary by design, recompute the treedef per call or register a custom pytree node that treats length as dynamic aux data.
4. Add assertions on len() before calling JAX APIs to fail early with a clear message.

Example fix

# before
batch = [x1, x2]                  # treedef expects 3 elements
# after
batch = [x1, x2, x3]              # or pad: batch += [dummy]*(3-len(batch))

expected_len = 3
if len(batch_list) != expected_len:
    batch_list = (batch_list + [pad] * expected_len)[:expected_len]

def list_arity_ok(lst: list, arity: int) -> bool:
    return type(lst) is list and len(lst) == arity

try:
    treedef.flatten_up_to(obj)
except ValueError as e:
    if 'List arity mismatch' in str(e):
        # resize then retry once
        raise
    raise

When it happens

Trigger: Flattening an object against a prefix treedef whose list node has arity N while the runtime list has a different number of elements; typical in jit-compiled function reuse, tree_map over trees where one side's list length changed, or replaying stored treedefs on batch data of different size.

Common situations: Variable batch sizes or sequence lengths stored in lists and passed to a jitted function whose cache expects the traced length, dynamic list building (append in loops) producing inconsistent lengths, editing datasets/config lists between runs.

Related errors

• Tuple arity mismatch: %d != %d; tuple: %s.
• Expected list, got %s.
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

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