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

Expected list, got %s.

Error message

Expected list, got %s.

What it means

JAX's PyTreeDef::FlattenUpTo found a kList node in the stored treedef, but the object at that position is not an exact built-in list (PyList_CheckExact). List subclasses and other sequences (tuples, arrays) are rejected because the treedef records exact types for container nodes.

Source

Thrown at jaxlib/pytree.cc:1020

          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(
              "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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass an exact list at that position: list(x).
2. Regenerate the treedef from the current object instead of reusing one captured before the container-type change.
3. Avoid converting containers between list/tuple between tracing and calling; keep one canonical type.
4. Use type-stable construction helpers (always list comprehensions, not conditional tuple/list).

Example fix

# before
obj = ((1, 2), (3, 4))            # inner was list in treedef
# after
obj = ([1, 2], [3, 4])            # exact lists

assert all(type(x) is list for x in my_lists), 'inner containers must be exact lists'

def is_exact_list_tree(obj, ref) -> bool:
    return jax.tree_util.tree_structure(obj).equals(
        jax.tree_util.tree_structure(ref)
    ) and all(type(l) is list for l in [obj] if isinstance(l, list))

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

When it happens

Trigger: Flattening against a prefix treedef where the node is a list but the runtime value is a tuple, numpy array, or a list subclass; commonly from jit cache reuse, tree structure replay, or converting containers between list and tuple during refactors.

Common situations: Converting list to tuple for hashing/immutability (e.g. caching keys) while the treedef still expects list, passing numpy arrays or generators where lists were traced, list subclasses from third-party libs (e.g. DataFrame rows-like subclasses).

Related errors

• Expected tuple, got %s.
• List arity mismatch: %d != %d; list: %s.
• Expected dict, got %s.
• Expected named tuple, got %s.
• numpy masked arrays are not supported as direct inputs to JA

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