{"record":{"id":"7d91f292b26d8513","repo":"jax-ml/jax","slug":"csc-tree-unflatten-invalid-aux-data","errorCode":null,"errorMessage":"CSC.tree_unflatten: invalid {aux_data=}","messagePattern":"CSC\\.tree_unflatten: invalid (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/csr.py","lineNumber":219,"sourceCode":"    data, other = promote_dtypes(self.data, other)\n    if other.ndim == 1:\n      return _csr_matvec(data, self.indices, self.indptr, other,\n                         shape=self.shape[::-1], transpose=True)\n    elif other.ndim == 2:\n      return _csr_matmat(data, self.indices, self.indptr, other,\n                         shape=self.shape[::-1], transpose=True)\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\"CSC.tree_unflatten: invalid {aux_data=}\")\n    obj.__dict__.update(**aux_data)\n    return obj\n\n\n#--------------------------------------------------------------------\n# csr_todense\n\ncsr_todense_p = core.Primitive('csr_todense')\n\ndef csr_todense(mat: CSR) -> Array:\n  \"\"\"Convert a CSR-format sparse matrix to a dense matrix.\n\n  Args:\n    mat : CSR matrix\n  Returns:\n    mat_dense: dense version of ``mat``\n  \"\"\"\n  return _csr_todense(mat.data, mat.indices, mat.indptr, shape=mat.shape)","sourceCodeStart":201,"sourceCodeEnd":237,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/csr.py#L201-L237","documentation":"JAX's CSC (compressed sparse column) array is a pytree; when it is unflattened (e.g. across jit/pmap boundaries, pickling, or tree_map), the auxiliary data must contain exactly the key 'shape'. If any other aux data is present, tree_unflatten rejects it as corrupt or incompatible.","triggerScenarios":"Reconstructing a jax.experimental.sparse.CSC from flattened children with aux_data whose keys are not exactly {'shape'}; typically from stale serialized/pickled objects from an older JAX version or manual tree manipulation.","commonSituations":"Unpickling CSC arrays saved by a different JAX version whose pytree layout changed; custom pytree plumbing that constructs aux_data dicts.","solutions":["Regenerate/re-save the sparse arrays with the current JAX version instead of unpickling old ones","If building aux_data manually, pass exactly {'shape': (n_rows, n_cols)}","Check that data, indices, indptr children order matches the CSC layout"],"exampleFix":"// before\naux = {'shape': (3, 4), 'nse': 5}  # extra key -> error\nCSC.tree_unflatten(aux, (data, indices, indptr))\n// after\naux = {'shape': (3, 4)}\nCSC.tree_unflatten(aux, (data, indices, indptr))","handlingStrategy":"validation","validationCode":"assert set(aux.keys()) == {'shape'}, f'bad aux keys: {aux.keys()}'","typeGuard":"def is_valid_csc_aux(aux) -> bool:\n    return isinstance(aux, dict) and set(aux.keys()) == {'shape'} and isinstance(aux['shape'], tuple)","tryCatchPattern":"try:\n    obj = CSC.tree_unflatten(aux, children)\nexcept ValueError as e:\n    if 'tree_unflatten' in str(e):\n        raise RuntimeError('Corrupt/incompatible CSC pytree data; re-save with current JAX') from e\n    raise","preventionTips":["Version-pin JAX across serialization and deserialization environments","Prefer sparse arrays' own save/load (e.g. store data/indices/indptr/shape explicitly) over pickling whole pytrees"],"tags":["jax","sparse","pytree","csc","serialization"],"backgroundTag":"pytree-unflatten-invalid-aux-data","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}