jax-ml/jax · error · ValueError

CSR.tree_unflatten: invalid {aux_data=}

Error message

CSR.tree_unflatten: invalid {aux_data=}

What it means

CSR is a pytree whose aux_data must be exactly {'shape'}. tree_unflatten raises ValueError on any other key set, guarding against pytrees flattened by different JAX versions or aux dicts built by hand. CSR and CSC share this contract (both only store shape as aux data).

Source

Thrown at jax/experimental/sparse/csr.py:144

      raise NotImplementedError("matmul between two sparse objects.")
    other = jnp.asarray(other)
    data, other = promote_dtypes(self.data, other)
    if other.ndim == 1:
      return _csr_matvec(data, self.indices, self.indptr, other, shape=self.shape)
    elif other.ndim == 2:
      return _csr_matmat(data, self.indices, self.indptr, other, shape=self.shape)
    else:
      raise NotImplementedError(f"matmul with object of shape {other.shape}")

  def tree_flatten(self):
    return (self.data, self.indices, self.indptr), {"shape": self.shape}

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    obj = object.__new__(cls)
    obj.data, obj.indices, obj.indptr = children
    if aux_data.keys() != {'shape'}:
      raise ValueError(f"CSR.tree_unflatten: invalid {aux_data=}")
    obj.__dict__.update(**aux_data)
    return obj


@tree_util.register_pytree_node_class
class CSC(JAXSparse):
  """Experimental CSC matrix implemented in JAX; API subject to change."""
  data: jax.Array
  indices: jax.Array
  indptr: jax.Array
  shape: tuple[int, int]  # pyrefly: ignore[bad-override]

  @property
  def nse(self) -> int:
    return self.data.size

  @property
  def dtype(self) -> np.dtype:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass aux_data={'shape': tuple} exactly when unflattening manually
2. Prefer pickle/jit round-trips of the live object over manual aux reconstruction
3. Migrate old serialized aux dicts to the current schema before loading

Example fix

# before
aux = {'shape': (4, 4), 'nse': 5}
obj = CSR.tree_unflatten(aux, children)  # ValueError

# after
aux = {'shape': (4, 4)}
obj = CSR.tree_unflatten(aux, children)

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

def csr_aux_valid(aux) -> bool:
    return set(aux.keys()) == {'shape'}

When it happens

Trigger: Reconstructing a CSR via tree_unflatten with aux_data containing extra keys (e.g. 'rows_sorted' copied from COO code) or missing 'shape'; deserializing checkpoints across JAX versions.

Common situations: Serialization/checkpointing of sparse pytrees; generic tree code that assumes all sparse classes share the same aux schema; version upgrades changing the aux layout.

Related errors

• COO.tree_unflatten: invalid {aux_data=}
• CSC.tree_unflatten: invalid {aux_data=}
• todense_transpose for {type(obj)}
• CSR must have ndim=2; got {shape=}
• matmul between two sparse objects.

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