jax-ml/jax · error · ValueError

COO.tree_unflatten: invalid {aux_data=}

Error message

COO.tree_unflatten: invalid {aux_data=}

What it means

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.

Source

Thrown at jax/experimental/sparse/coo.py:162

  def todense(self) -> Array:
    return coo_todense(self)

  def transpose(self, axes: tuple[int, ...] | None = None) -> COO:
    if axes is not None:
      raise NotImplementedError("axes argument to transpose()")
    return COO((self.data, self.col, self.row), shape=self.shape[::-1],
               rows_sorted=self._cols_sorted, cols_sorted=self._rows_sorted)

  def tree_flatten(self) -> tuple[tuple[Array, Array, Array], dict[str, Any]]:
    return (self.data, self.row, self.col), self._info._asdict()

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    obj = object.__new__(cls)
    obj.data, obj.row, obj.col = children
    if aux_data.keys() != {'shape', 'rows_sorted', 'cols_sorted'}:
      raise ValueError(f"COO.tree_unflatten: invalid {aux_data=}")
    obj.shape = aux_data['shape']
    obj._rows_sorted = aux_data['rows_sorted']
    obj._cols_sorted = aux_data['cols_sorted']
    return obj

  def __matmul__(self, other: ArrayLike) -> Array:
    if isinstance(other, JAXSparse):
      raise NotImplementedError("matmul between two sparse objects.")
    other = jnp.asarray(other)
    data, other = promote_dtypes(self.data, other)
    self_promoted = COO((data, self.row, self.col), **self._info._asdict())
    if other.ndim == 1:
      return coo_matvec(self_promoted, other)
    elif other.ndim == 2:
      return coo_matmat(self_promoted, other)
    else:
      raise NotImplementedError(f"matmul with object of shape {other.shape}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Regenerate the aux_data from a live COO (coo._info._asdict()) rather than storing a hand-built dict
2. Re-serialize/re-checkpoint using the same JAX version that will read it
3. If loading old checkpoints, migrate aux dicts to the expected key set {'shape','rows_sorted','cols_sorted'} before unflattening

Example fix

# before
aux = {'shape': (4, 4), 'rows_sorted': True}  # missing cols_sorted
obj = COO.tree_unflatten(aux, (data, row, col))  # ValueError

# after
aux = {'shape': (4, 4), 'rows_sorted': True, 'cols_sorted': True}
obj = COO.tree_unflatten(aux, (data, row, col))

assert set(aux_data) == {'shape', 'rows_sorted', 'cols_sorted'}, aux_data

def coo_aux_valid(aux) -> bool:
    return set(aux.keys()) == {'shape', 'rows_sorted', 'cols_sorted'}

When it happens

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

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

Related errors

• CSR.tree_unflatten: invalid {aux_data=}
• CSC.tree_unflatten: invalid {aux_data=}
• Unsupported shape: {shape}
• COO must have ndim=2; got {shape=}
• axes argument to transpose()

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