{"record":{"id":"6e7569ec1f53114d","repo":"jax-ml/jax","slug":"expected-dict-got-s","errorCode":null,"errorMessage":"Expected dict, got %s.","messagePattern":"Expected dict, got (.+?)\\.","errorType":"validation","errorClass":"std::invalid_argument","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1038,"sourceCode":"          throw std::invalid_argument(\n              absl::StrFormat(\"Expected list, got %s.\",\n                              nb::cast<std::string_view>(nb::repr(object))));\n        }\n        nb::list list = nb::borrow<nb::list>(object);\n        if (list.size() != node.arity) {\n          throw std::invalid_argument(absl::StrFormat(\n              \"List arity mismatch: %d != %d; list: %s.\", list.size(),\n              node.arity, nb::cast<std::string_view>(nb::repr(object))));\n        }\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;","sourceCodeStart":1020,"sourceCodeEnd":1056,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L1020-L1056","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Convert to a plain dict at that position: dict(x) (or {**x}).","Re-trace/re-derive the treedef after library upgrades changed container types (e.g. flax>=0.5 FrozenDict migration).","Pin library versions consistently across environments so the same dict type is produced.","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."],"exampleFix":"# before\nparams = FrozenDict(...)          # treedef node is plain dict\n# after\nparams = dict(FrozenDict(...))    # or upgrade flax and re-trace","handlingStrategy":"type-guard","validationCode":"assert type(params) is dict, f'need plain dict, got {type(params).__name__}'; params = dict(params)","typeGuard":"def is_exact_dict(x) -> bool:\n    return type(x) is dict","tryCatchPattern":"try:\n    treedef.flatten_up_to(obj)\nexcept (ValueError, TypeError) as e:\n    if 'Expected dict' in str(e):\n        obj = jax.tree_util.tree_map(\n            lambda x: dict(x) if isinstance(x, dict) and type(x) is not dict else x,\n            obj,\n        )\n    else:\n        raise","preventionTips":["Normalize Mapping inputs with dict(x) at API boundaries.","Pin jax/flax versions to keep dict types (FrozenDict vs dict) stable.","Register custom Mapping classes as pytree nodes instead of relying on dict compat."],"tags":["jax","pytree","type-mismatch","dict"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}