jax-ml/jax · error · NotImplementedError

matmul between two sparse objects.

Error message

matmul between two sparse objects.

What it means

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.

Source

Thrown at jax/experimental/sparse/coo.py:170

               rows_sorted=self._cols_sorted, cols_sorted=self._rows_sorted)

  def tree_flatten(self) -> tuple[tuple[Array, Array, Array], dict[str, Any]]:
    return (self.data, self.row, self.col), self._info._asdict()

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    obj = object.__new__(cls)
    obj.data, obj.row, obj.col = children
    if aux_data.keys() != {'shape', 'rows_sorted', 'cols_sorted'}:
      raise ValueError(f"COO.tree_unflatten: invalid {aux_data=}")
    obj.shape = aux_data['shape']
    obj._rows_sorted = aux_data['rows_sorted']
    obj._cols_sorted = aux_data['cols_sorted']
    return obj

  def __matmul__(self, other: ArrayLike) -> Array:
    if isinstance(other, JAXSparse):
      raise NotImplementedError("matmul between two sparse objects.")
    other = jnp.asarray(other)
    data, other = promote_dtypes(self.data, other)
    self_promoted = COO((data, self.row, self.col), **self._info._asdict())
    if other.ndim == 1:
      return coo_matvec(self_promoted, other)
    elif other.ndim == 2:
      return coo_matmat(self_promoted, other)
    else:
      raise NotImplementedError(f"matmul with object of shape {other.shape}")

#--------------------------------------------------------------------
# coo_todense

coo_todense_p = core.Primitive('coo_todense')

def coo_todense(mat: COO) -> Array:
  """Convert a COO-format sparse matrix to a dense matrix.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Densify one operand: coo_a @ coo_b.todense()
2. Use BCOO: bcoo.bcoo_matmul(BCOO.from_coo(a), BCOO.from_coo(b)) (or sparse.sparsify(jnp.matmul))
3. For scipy-equivalent semantics, drop to scipy.sparse for the sparse-sparse product and convert back

Example fix

# before
c = coo_a @ coo_b  # NotImplementedError

# after
c = bcoo.bcoo_matmul(bcoo.BCOO.fromdense(coo_a.todense()),
                     bcoo.BCOO.fromdense(coo_b.todense()))

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

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

try:
    c = a @ b
except NotImplementedError:
    c = bcoo.bcoo_matmul(bcoo.BCOO.fromdense(a.todense()),
                         bcoo.BCOO.fromdense(b.todense()))

When it happens

Trigger: coo_a @ coo_b, coo_a @ csr_b, or any '@' where the right side isinstance of JAXSparse (COO, CSR, CSC, BCOO, BCSR).

Common situations: Composing sparse factor matrices (e.g. sparse diag @ sparse matrix); porting scipy code where scipy handles sparse@sparse natively.

Related errors

• axes argument to transpose()
• matmul with object of shape {other.shape}
• matmul between two sparse objects.
• Transforms not supported for matmul_acc_lhs.
• Transforms not supported for matmul_pop.

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