{"record":{"id":"19deb0a76062a9a5","repo":"jax-ml/jax","slug":"csr-tree-unflatten-invalid-aux-data","errorCode":null,"errorMessage":"CSR.tree_unflatten: invalid {aux_data=}","messagePattern":"CSR\\.tree_unflatten: invalid (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/csr.py","lineNumber":144,"sourceCode":"      raise NotImplementedError(\"matmul between two sparse objects.\")\n    other = jnp.asarray(other)\n    data, other = promote_dtypes(self.data, other)\n    if other.ndim == 1:\n      return _csr_matvec(data, self.indices, self.indptr, other, shape=self.shape)\n    elif other.ndim == 2:\n      return _csr_matmat(data, self.indices, self.indptr, other, shape=self.shape)\n    else:\n      raise NotImplementedError(f\"matmul with object of shape {other.shape}\")\n\n  def tree_flatten(self):\n    return (self.data, self.indices, self.indptr), {\"shape\": self.shape}\n\n  @classmethod\n  def tree_unflatten(cls, aux_data, children):\n    obj = object.__new__(cls)\n    obj.data, obj.indices, obj.indptr = children\n    if aux_data.keys() != {'shape'}:\n      raise ValueError(f\"CSR.tree_unflatten: invalid {aux_data=}\")\n    obj.__dict__.update(**aux_data)\n    return obj\n\n\n@tree_util.register_pytree_node_class\nclass CSC(JAXSparse):\n  \"\"\"Experimental CSC matrix implemented in JAX; API subject to change.\"\"\"\n  data: jax.Array\n  indices: jax.Array\n  indptr: jax.Array\n  shape: tuple[int, int]  # pyrefly: ignore[bad-override]\n\n  @property\n  def nse(self) -> int:\n    return self.data.size\n\n  @property\n  def dtype(self) -> np.dtype:","sourceCodeStart":126,"sourceCodeEnd":162,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/csr.py#L126-L162","documentation":"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).","triggerScenarios":"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.","commonSituations":"Serialization/checkpointing of sparse pytrees; generic tree code that assumes all sparse classes share the same aux schema; version upgrades changing the aux layout.","solutions":["Pass aux_data={'shape': tuple} exactly when unflattening manually","Prefer pickle/jit round-trips of the live object over manual aux reconstruction","Migrate old serialized aux dicts to the current schema before loading"],"exampleFix":"# before\naux = {'shape': (4, 4), 'nse': 5}\nobj = CSR.tree_unflatten(aux, children)  # ValueError\n\n# after\naux = {'shape': (4, 4)}\nobj = CSR.tree_unflatten(aux, children)","handlingStrategy":"validation","validationCode":"assert set(aux_data) == {'shape'}, aux_data","typeGuard":"def csr_aux_valid(aux) -> bool:\n    return set(aux.keys()) == {'shape'}","tryCatchPattern":null,"preventionTips":["Use aux={'shape': shape} exactly for CSR/CSC unflattening","Rely on pickle/round-trip rather than manual aux reconstruction"],"tags":["jax","sparse","csr","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"}