{"record":{"id":"124c5948326a397d","repo":"jax-ml/jax","slug":"could-not-find-type-s","errorCode":null,"errorMessage":"Could not find type: %s.","messagePattern":"Could not find type: (.+?)\\.","errorType":"exception","errorClass":"logic_error","httpStatus":null,"severity":"error","filePath":"jaxlib/pytree.cc","lineNumber":1636,"sourceCode":"  node.num_nodes = result->traversal_.size();\n  if (node_data == std::nullopt) {\n    node.kind = PyTreeKind::kLeaf;\n    ++node.num_leaves;\n    return result;\n  }\n  int is_nt = PyObject_IsSubclass(node_data->first.ptr(),\n                                  reinterpret_cast<PyObject*>(&PyTuple_Type));\n  if (is_nt == -1) {\n    throw nb::python_error();\n  }\n  if (is_nt != 0 && nb::hasattr(node_data->first, \"_fields\")) {\n    node.kind = PyTreeKind::kNamedTuple;\n    node.node_data = node_data->first;\n    return result;\n  }\n  auto* registration = result->registry()->Lookup(node_data->first);\n  if (registration == nullptr) {\n    throw std::logic_error(absl::StrFormat(\n        \"Could not find type: %s.\",\n        nb::cast<absl::string_view>(nb::repr(node_data->first))));\n  }\n  node.kind = registration->kind;\n  if (node.kind == PyTreeKind::kCustom || node.kind == PyTreeKind::kDataclass) {\n    node.custom = registration;\n    node.node_data = node_data->second;\n  } else if (node.kind == PyTreeKind::kNamedTuple) {\n    node.node_data = node_data->first;\n  } else if (node.kind == PyTreeKind::kDict) {\n    node.sorted_dict_keys =\n        nb::cast<std::vector<nb::object>>(node_data->second);\n  }\n  return result;\n}\n\nint PyTreeDef::Node::tp_traverse(visitproc visit, void* arg) const {\n  Py_VISIT(node_data.ptr());","sourceCodeStart":1618,"sourceCodeEnd":1654,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jaxlib/pytree.cc#L1618-L1654","documentation":"Raised when reconstructing a PyTreeDef from node data whose type is not registered in the pytree registry. The node's type must have been registered via jax.tree_util.register_pytree_node (or pytree registration) for deserialization/registration to succeed.","triggerScenarios":"Deserializing or unpickling a treedef containing a custom class whose pytree registration hasn't happened yet in the current process; the registry lookup by type/repr fails and a std::logic_error is thrown.","commonSituations":"Unpickling JAX objects (e.g. across process boundaries with multiprocessing/cloudpickle) where the custom container class isn't imported or registered in the child; version skew where registration names changed; registering a different object under the same name.","solutions":["Ensure the module that calls jax.tree_util.register_pytree_node(...) for the custom type is imported before unpickling/deserializing","Import the defining library (e.g. flax, custom dataclasses) in the worker process before receiving pickled treedefs","Check registration identity: same class object, not a redefined class with the same name","Pin matching jax/jaxlib versions on producer and consumer"],"exampleFix":"# before\n# in worker: result = pickle.loads(payload)  # custom class not imported\n\n# after\nimport mypkg.containers  # performs register_pytree_node at import time\nresult = pickle.loads(payload)","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\n# ensure the custom type is registered before unpickling/deserializing\nimport mypkg.containers  # module registers pytree at import\nassert any(getattr(t, '__name__', '') == 'MyContainer'\n           for t in vars(mypkg.containers) if isinstance(t, type))","typeGuard":"def is_registered_pytree(cls) -> bool:\n    import jax.tree_util as jtu\n    try:\n        jtu.tree_structure(cls.__new__(cls) if hasattr(cls, '__new__') else None)\n        return True\n    except Exception:\n        return False","tryCatchPattern":"try:\n    obj = pickle.loads(payload)\nexcept Exception as e:\n    if 'Could not find type' in str(e):\n        import mypkg.containers  # register, then retry once\n        obj = pickle.loads(payload)\n    else:\n        raise","preventionTips":["Import registration modules at process start in workers (multiprocessing initializer)","Never reload/redefine registered pytree classes","Keep the same class identity across processes (no dynamic redefinition)"],"tags":["jax","pytree","pickle","registry"],"backgroundTag":"missing-type-registration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}