jax-ml/jax · error · NotImplementedError

matmul between two sparse objects.

Error message

matmul between two sparse objects.

What it means

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.

Source

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

    k = _const(idx, k)
    col = lax.add(idx, lax.cond(k <= 0, lambda: zero, lambda: k))
    indices = col.astype(index_dtype)
    # TODO(jakevdp): this can be done more efficiently.
    row = lax.sub(idx, lax.cond(k >= 0, lambda: zero, lambda: k))
    indptr = jnp.zeros(N + 1, dtype=index_dtype).at[1:].set(
        jnp.cumsum(jnp.bincount(row, length=N).astype(index_dtype)))
    return cls((data, indices, indptr), shape=(N, M))

  def todense(self):
    return csr_todense(self)

  def transpose(self, axes=None):
    assert axes is None
    return CSC((self.data, self.indices, self.indptr), shape=self.shape[::-1])

  def __matmul__(self, other):
    if isinstance(other, JAXSparse):
      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=}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Densify one side: csr_a @ csr_b.todense()
2. Use bcoo.bcoo_matmul / sparse.sparsify(jnp.matmul) with BCOO operands
3. Use BCSR and bcsr matmul routines for batched 2D sparse-sparse cases

Example fix

# before
c = csr_a @ csr_b  # NotImplementedError

# after
c = bcoo.bcoo_matmul(bcoo.BCOO.fromdense(csr_a.todense()),
                     bcoo.BCOO.fromdense(csr_b.todense()))

from jax.experimental.sparse import JAXSparse
assert not isinstance(rhs, JAXSparse), 'CSR matmul requires a dense rhs'

def csr_matmul_rhs_ok(other) -> bool:
    from jax.experimental.sparse import JAXSparse
    return not isinstance(other, JAXSparse)

try:
    c = a @ b
except NotImplementedError:
    c = a @ b.todense()

When it happens

Trigger: csr_a @ csr_b, csr_a @ coo_b, or any '@' where the right operand isinstance(other, JAXSparse).

Common situations: Porting scipy.sparse code where sparse @ sparse is routine; composing sparse linear operators (L = D @ A with both sparse).

Related errors

• matmul between two sparse objects.
• matmul with object of shape {other.shape}
• Transforms not supported for matmul_acc_lhs.
• Transforms not supported for matmul_pop.
• Sparse metadata format not implemented for {operand_dtype=}

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