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

Dict key mismatch; expected keys

Error message

Dict key mismatch; expected keys: %s; present keys: %s.

What it means

During FlattenUpTo, JAX compared the sorted key sets of the runtime dict against the keys recorded in the treedef node and found them unequal. Pytree dicts are matched by exact key sets, so an inserted, removed, or renamed key makes the structure mismatch. The message prints both expected and present key lists.

Solutions

Align the dict keys to the expected set printed in the message (add missing, drop extra: {k: v for k, v in d.items() if k in expected}).
Recompute the treedef from the current object (re-jit / fresh tree_structure) after key changes.
Normalize dicts through a fixed schema/keys function before any JAX API call.
For optional keys, use a stable placeholder (e.g. None) instead of adding/removing keys.

Example fix

# before
obj = {'a': 1, 'b': 2}            # treedef expects {'a','c'}
# after
obj = {'a': 1, 'c': obj.pop('b')} # keys match treedef

Defensive patterns

Strategy: validation

Validate before calling

expected_keys = {'a', 'b', 'c'}
assert set(my_dict) == expected_keys, f'key drift: extra={set(my_dict)-expected_keys}, missing={expected_keys-set(my_dict)}'

Type guard

def keys_match(d: dict, expected: set[str]) -> bool:
    return type(d) is dict and set(d) == expected

Try / catch

try:
    treedef.flatten_up_to(obj)
except ValueError as e:
    if 'Dict key mismatch' in str(e):
        expected = set(eval(e.split('expected keys: ')[1].split(';')[0]))  # or re-derive
        obj = {k: v for k, v in obj.items() if k in expected}
    else:
        raise

Prevention

Use fixed schema/validation (pydantic or dataclass) for config dicts before JAX calls.
Avoid add/remove of dict keys between trace and call; use None placeholders.
Diff keys in tests: assert set(new_config) == set(golden_config).

When it happens

Trigger: Flattening an object whose dict has different keys (added/removed/renamed) than the dict captured in the treedef; e.g. tree_unflatten followed by key edits then re-flattening against the old treedef, jit argument structure drift, applying a treedef from one model config to another with extra hyperparameter keys.

Common situations: Adding a new config/param key between experiments while reusing compiled functions or stored treedefs, renaming dict keys during refactors, datasets with per-sample metadata dicts of varying keys, flax params dicts differing across module versions.

Related errors

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

Appendix: source

Thrown at jaxlib/pytree.cc:1047

        }
        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;
      }

      case PyTreeKind::kNamedTuple: {
        if (!nb::isinstance<nb::tuple>(object) ||
            !nb::hasattr(object, "_fields")) {
          throw std::invalid_argument(
              absl::StrFormat("Expected named tuple, got %s.",
                              nb::cast<std::string_view>(nb::repr(object))));
        }

View on GitHub (pinned to 1e1c6a8fc0)