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

Expected dict, got %s.

Error message

Expected dict, got %s.

What it means

JAX's FlattenUpTo hit a kDict node in the treedef, but the object at that position is not an exact built-in dict (PyDict_CheckExact). dict subclasses (OrderedDict, defaultdict, custom Mapping subclasses) and other mappings are rejected at this node because the treedef stores exact container types.

Source

Thrown at jaxlib/pytree.cc:1038

          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)) {
          // 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;

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to a plain dict at that position: dict(x) (or {**x}).
2. Re-trace/re-derive the treedef after library upgrades changed container types (e.g. flax>=0.5 FrozenDict migration).
3. Pin library versions consistently across environments so the same dict type is produced.
4. Register custom Mapping classes as pytree nodes via jax.tree_util.register_pytree_node so their structure is handled by aux_data, not exact dict checks.

Example fix

# before
params = FrozenDict(...)          # treedef node is plain dict
# after
params = dict(FrozenDict(...))    # or upgrade flax and re-trace

assert type(params) is dict, f'need plain dict, got {type(params).__name__}'; params = dict(params)

def is_exact_dict(x) -> bool:
    return type(x) is dict

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

When it happens

Trigger: Flattening against a treedef that captured a plain dict while the runtime value is an OrderedDict/defaultdict/CustomDict or a non-dict; happens when reusing treedefs across code that changed dict flavors, or passing params as flax FrozenDict/CustomDict where a plain dict was traced (or vice versa after flax version changes).

Common situations: Upgrading Flax (FrozenDict removed/changed) so param pyramids are different dict types, replacing dicts with Mapping dataclasses, mixing to_dict() outputs across library versions, reusing pickled treedefs.

Related errors

• Expected tuple, got %s.
• Expected list, got %s.
• Dict key mismatch; expected keys: %s; present keys: %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/6e7569ec1f53114d. Report an issue: GitHub.