{"record":{"id":"85c9e3ca933b3d53","repo":"jax-ml/jax","slug":"coo-tree-unflatten-invalid-aux-data","errorCode":null,"errorMessage":"COO.tree_unflatten: invalid {aux_data=}","messagePattern":"COO\\.tree_unflatten: invalid (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/coo.py","lineNumber":162,"sourceCode":"\n  def todense(self) -> Array:\n    return coo_todense(self)\n\n  def transpose(self, axes: tuple[int, ...] | None = None) -> COO:\n    if axes is not None:\n      raise NotImplementedError(\"axes argument to transpose()\")\n    return COO((self.data, self.col, self.row), shape=self.shape[::-1],\n               rows_sorted=self._cols_sorted, cols_sorted=self._rows_sorted)\n\n  def tree_flatten(self) -> tuple[tuple[Array, Array, Array], dict[str, Any]]:\n    return (self.data, self.row, self.col), self._info._asdict()\n\n  @classmethod\n  def tree_unflatten(cls, aux_data, children):\n    obj = object.__new__(cls)\n    obj.data, obj.row, obj.col = children\n    if aux_data.keys() != {'shape', 'rows_sorted', 'cols_sorted'}:\n      raise ValueError(f\"COO.tree_unflatten: invalid {aux_data=}\")\n    obj.shape = aux_data['shape']\n    obj._rows_sorted = aux_data['rows_sorted']\n    obj._cols_sorted = aux_data['cols_sorted']\n    return obj\n\n  def __matmul__(self, other: ArrayLike) -> Array:\n    if isinstance(other, JAXSparse):\n      raise NotImplementedError(\"matmul between two sparse objects.\")\n    other = jnp.asarray(other)\n    data, other = promote_dtypes(self.data, other)\n    self_promoted = COO((data, self.row, self.col), **self._info._asdict())\n    if other.ndim == 1:\n      return coo_matvec(self_promoted, other)\n    elif other.ndim == 2:\n      return coo_matmat(self_promoted, other)\n    else:\n      raise NotImplementedError(f\"matmul with object of shape {other.shape}\")\n","sourceCodeStart":144,"sourceCodeEnd":180,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/coo.py#L144-L180","documentation":"COO is a pytree class; tree_unflatten reconstructs it from children (data, row, col) and an aux_data dict that must contain exactly the keys 'shape', 'rows_sorted', 'cols_sorted'. If the aux dict has different keys (added, missing, or renamed), JAX raises ValueError — this guards against pytree serialization/version mismatches and corrupted transforms.","triggerScenarios":"Round-tripping a COO through jax.tree_util tree flatten/unflatten or serialization where aux_data was produced by a different JAX version or manually constructed with wrong keys; custom pytree code that copies _info._asdict() and mutates it.","commonSituations":"Pickling/serializing COO pytrees across JAX versions whose _COOInfo fields changed; custom checkpointing code that reconstructs aux dicts; flax/optimystic-style transformations touching aux data.","solutions":["Regenerate the aux_data from a live COO (coo._info._asdict()) rather than storing a hand-built dict","Re-serialize/re-checkpoint using the same JAX version that will read it","If loading old checkpoints, migrate aux dicts to the expected key set {'shape','rows_sorted','cols_sorted'} before unflattening"],"exampleFix":"# before\naux = {'shape': (4, 4), 'rows_sorted': True}  # missing cols_sorted\nobj = COO.tree_unflatten(aux, (data, row, col))  # ValueError\n\n# after\naux = {'shape': (4, 4), 'rows_sorted': True, 'cols_sorted': True}\nobj = COO.tree_unflatten(aux, (data, row, col))","handlingStrategy":"validation","validationCode":"assert set(aux_data) == {'shape', 'rows_sorted', 'cols_sorted'}, aux_data","typeGuard":"def coo_aux_valid(aux) -> bool:\n    return set(aux.keys()) == {'shape', 'rows_sorted', 'cols_sorted'}","tryCatchPattern":null,"preventionTips":["Derive aux dicts via coo._info._asdict() instead of building by hand","Re-checkpoint sparse pytrees when upgrading JAX versions"],"tags":["jax","sparse","coo","pytree","serialization"],"backgroundTag":"pytree-aux-data-invalid","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}