{"record":{"id":"bb0d519e780f0719","repo":"jax-ml/jax","slug":"dict-key-mismatch-expected-keys-s-present-keys","errorCode":null,"errorMessage":"Dict key mismatch; expected keys: %s; present keys: %s.","messagePattern":"Dict key mismatch; expected keys: (.+?); present keys: (.+?)\\.","errorType":"validation","errorClass":"std::invalid_argument","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1047,"sourceCode":"        }\n        for (nb::handle entry : list) {\n          agenda.push_back(nb::borrow<nb::object>(entry));\n        }\n        break;\n      }\n\n      case PyTreeKind::kDict: {\n        if (!PyDict_CheckExact(object.ptr())) {\n          throw std::invalid_argument(\n              absl::StrFormat(\"Expected dict, got %s.\",\n                              nb::cast<std::string_view>(nb::repr(object))));\n        }\n        nb::dict dict = nb::borrow<nb::dict>(object);\n        std::vector<nb::object> keys = GetSortedPyDictKeys(dict.ptr());\n        if (!IsSortedPyDictKeysEqual(keys, node.sorted_dict_keys)) {\n          // Convert to a nb::list for nb::repr to avoid having to stringify a\n          // vector. This is error path so it is fine to pay conversion cost.\n          throw std::invalid_argument(absl::StrFormat(\n              \"Dict key mismatch; expected keys: %s; present keys: %s.\",\n              nb::cast<std::string_view>(\n                  nb::repr(nb::cast(node.sorted_dict_keys))),\n              nb::cast<std::string_view>(nb::repr(nb::cast(keys)))));\n        }\n        for (nb::handle key : keys) {\n          agenda.push_back(dict[key]);\n        }\n        break;\n      }\n\n      case PyTreeKind::kNamedTuple: {\n        if (!nb::isinstance<nb::tuple>(object) ||\n            !nb::hasattr(object, \"_fields\")) {\n          throw std::invalid_argument(\n              absl::StrFormat(\"Expected named tuple, got %s.\",\n                              nb::cast<std::string_view>(nb::repr(object))));\n        }","sourceCodeStart":1029,"sourceCodeEnd":1065,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L1029-L1065","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"# before\nobj = {'a': 1, 'b': 2}            # treedef expects {'a','c'}\n# after\nobj = {'a': 1, 'c': obj.pop('b')} # keys match treedef","handlingStrategy":"validation","validationCode":"expected_keys = {'a', 'b', 'c'}\nassert set(my_dict) == expected_keys, f'key drift: extra={set(my_dict)-expected_keys}, missing={expected_keys-set(my_dict)}'","typeGuard":"def keys_match(d: dict, expected: set[str]) -> bool:\n    return type(d) is dict and set(d) == expected","tryCatchPattern":"try:\n    treedef.flatten_up_to(obj)\nexcept ValueError as e:\n    if 'Dict key mismatch' in str(e):\n        expected = set(eval(e.split('expected keys: ')[1].split(';')[0]))  # or re-derive\n        obj = {k: v for k, v in obj.items() if k in expected}\n    else:\n        raise","preventionTips":["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)."],"tags":["jax","pytree","dict","key-mismatch"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}