jax-ml/jax · error · logic_error

Could not find type: %s.

Error message

Could not find type: %s.

What it means

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.

Source

Thrown at jaxlib/pytree.cc:1636

  node.num_nodes = result->traversal_.size();
  if (node_data == std::nullopt) {
    node.kind = PyTreeKind::kLeaf;
    ++node.num_leaves;
    return result;
  }
  int is_nt = PyObject_IsSubclass(node_data->first.ptr(),
                                  reinterpret_cast<PyObject*>(&PyTuple_Type));
  if (is_nt == -1) {
    throw nb::python_error();
  }
  if (is_nt != 0 && nb::hasattr(node_data->first, "_fields")) {
    node.kind = PyTreeKind::kNamedTuple;
    node.node_data = node_data->first;
    return result;
  }
  auto* registration = result->registry()->Lookup(node_data->first);
  if (registration == nullptr) {
    throw std::logic_error(absl::StrFormat(
        "Could not find type: %s.",
        nb::cast<absl::string_view>(nb::repr(node_data->first))));
  }
  node.kind = registration->kind;
  if (node.kind == PyTreeKind::kCustom || node.kind == PyTreeKind::kDataclass) {
    node.custom = registration;
    node.node_data = node_data->second;
  } else if (node.kind == PyTreeKind::kNamedTuple) {
    node.node_data = node_data->first;
  } else if (node.kind == PyTreeKind::kDict) {
    node.sorted_dict_keys =
        nb::cast<std::vector<nb::object>>(node_data->second);
  }
  return result;
}

int PyTreeDef::Node::tp_traverse(visitproc visit, void* arg) const {
  Py_VISIT(node_data.ptr());

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the module that calls jax.tree_util.register_pytree_node(...) for the custom type is imported before unpickling/deserializing
2. Import the defining library (e.g. flax, custom dataclasses) in the worker process before receiving pickled treedefs
3. Check registration identity: same class object, not a redefined class with the same name
4. Pin matching jax/jaxlib versions on producer and consumer

Example fix

# before
# in worker: result = pickle.loads(payload)  # custom class not imported

# after
import mypkg.containers  # performs register_pytree_node at import time
result = pickle.loads(payload)

import jax.tree_util as jtu
# ensure the custom type is registered before unpickling/deserializing
import mypkg.containers  # module registers pytree at import
assert any(getattr(t, '__name__', '') == 'MyContainer'
           for t in vars(mypkg.containers) if isinstance(t, type))

def is_registered_pytree(cls) -> bool:
    import jax.tree_util as jtu
    try:
        jtu.tree_structure(cls.__new__(cls) if hasattr(cls, '__new__') else None)
        return True
    except Exception:
        return False

try:
    obj = pickle.loads(payload)
except Exception as e:
    if 'Could not find type' in str(e):
        import mypkg.containers  # register, then retry once
        obj = pickle.loads(payload)
    else:
        raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• The names should be exclusive and should not intersect in `n
• the rematted computation's closure contains a mutable array
• Effects not supported in partial-eval of `checkpoint`/`remat
• stop_gradient only works on valid JAX arrays, but input argu
• {self.__class__.__name__} has no attribute {name}

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