{"record":{"id":"24d323cb2c6684ea","repo":"jax-ml/jax","slug":"matmul-between-two-sparse-objects-24d323","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/csr.py","lineNumber":126,"sourceCode":"    k = _const(idx, k)\n    col = lax.add(idx, lax.cond(k <= 0, lambda: zero, lambda: k))\n    indices = col.astype(index_dtype)\n    # TODO(jakevdp): this can be done more efficiently.\n    row = lax.sub(idx, lax.cond(k >= 0, lambda: zero, lambda: k))\n    indptr = jnp.zeros(N + 1, dtype=index_dtype).at[1:].set(\n        jnp.cumsum(jnp.bincount(row, length=N).astype(index_dtype)))\n    return cls((data, indices, indptr), shape=(N, M))\n\n  def todense(self):\n    return csr_todense(self)\n\n  def transpose(self, axes=None):\n    assert axes is None\n    return CSC((self.data, self.indices, self.indptr), shape=self.shape[::-1])\n\n  def __matmul__(self, other):\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    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=}\")","sourceCodeStart":108,"sourceCodeEnd":144,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/csr.py#L108-L144","documentation":"The legacy CSR class implements matmul only for a sparse matrix times a dense array (ndim 1 or 2). When the right operand is another JAXSparse instance (CSR, CSC, COO, BCOO, BCSR), CSR.__matmul__ raises NotImplementedError. Use BCOO/BCSR APIs for sparse-sparse products.","triggerScenarios":"csr_a @ csr_b, csr_a @ coo_b, or any '@' where the right operand isinstance(other, JAXSparse).","commonSituations":"Porting scipy.sparse code where sparse @ sparse is routine; composing sparse linear operators (L = D @ A with both sparse).","solutions":["Densify one side: csr_a @ csr_b.todense()","Use bcoo.bcoo_matmul / sparse.sparsify(jnp.matmul) with BCOO operands","Use BCSR and bcsr matmul routines for batched 2D sparse-sparse cases"],"exampleFix":"# before\nc = csr_a @ csr_b  # NotImplementedError\n\n# after\nc = bcoo.bcoo_matmul(bcoo.BCOO.fromdense(csr_a.todense()),\n                     bcoo.BCOO.fromdense(csr_b.todense()))","handlingStrategy":"type-guard","validationCode":"from jax.experimental.sparse import JAXSparse\nassert not isinstance(rhs, JAXSparse), 'CSR matmul requires a dense rhs'","typeGuard":"def csr_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 = a @ b.todense()","preventionTips":["Densify one operand for sparse-sparse products","Use bcoo_matmul / bcsr routines for sparse-sparse"],"tags":["jax","sparse","csr","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"}