{"record":{"id":"b24358fc87e33b7a","repo":"jax-ml/jax","slug":"matmul-between-two-sparse-objects","errorCode":null,"errorMessage":"matmul between two sparse objects.","messagePattern":"matmul between two sparse objects\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/coo.py","lineNumber":170,"sourceCode":"               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\n#--------------------------------------------------------------------\n# coo_todense\n\ncoo_todense_p = core.Primitive('coo_todense')\n\ndef coo_todense(mat: COO) -> Array:\n  \"\"\"Convert a COO-format sparse matrix to a dense matrix.\n","sourceCodeStart":152,"sourceCodeEnd":188,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/coo.py#L152-L188","documentation":"The legacy COO format implements matmul only between a sparse matrix and a dense array (ndim 1 or 2). Multiplying two sparse objects is not supported, so COO.__matmul__ raises NotImplementedError when the right operand is also a JAXSparse instance. Use the batched BCOO API for sparse-sparse products.","triggerScenarios":"coo_a @ coo_b, coo_a @ csr_b, or any '@' where the right side isinstance of JAXSparse (COO, CSR, CSC, BCOO, BCSR).","commonSituations":"Composing sparse factor matrices (e.g. sparse diag @ sparse matrix); porting scipy code where scipy handles sparse@sparse natively.","solutions":["Densify one operand: coo_a @ coo_b.todense()","Use BCOO: bcoo.bcoo_matmul(BCOO.from_coo(a), BCOO.from_coo(b)) (or sparse.sparsify(jnp.matmul))","For scipy-equivalent semantics, drop to scipy.sparse for the sparse-sparse product and convert back"],"exampleFix":"# before\nc = coo_a @ coo_b  # NotImplementedError\n\n# after\nc = bcoo.bcoo_matmul(bcoo.BCOO.fromdense(coo_a.todense()),\n                     bcoo.BCOO.fromdense(coo_b.todense()))","handlingStrategy":"type-guard","validationCode":"from jax.experimental.sparse import JAXSparse\nassert not isinstance(rhs, JAXSparse), 'COO matmul requires a dense rhs'","typeGuard":"def coo_matmul_rhs_ok(other) -> bool:\n    from jax.experimental.sparse import JAXSparse\n    return not isinstance(other, JAXSparse)","tryCatchPattern":"try:\n    c = a @ b\nexcept NotImplementedError:\n    c = bcoo.bcoo_matmul(bcoo.BCOO.fromdense(a.todense()),\n                         bcoo.BCOO.fromdense(b.todense()))","preventionTips":["Densify one operand for sparse-sparse products with legacy formats","Use bcoo_matmul when both operands must stay sparse"],"tags":["jax","sparse","coo","matmul","not-implemented"],"backgroundTag":"sparse-sparse-matmul-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}