jax-ml/jax · error · NotImplementedError

matmul with object of shape {other.shape}

Error message

matmul with object of shape {other.shape}

What it means

CSR.__matmul__ dispatches on the dense right operand's ndim: 1 → _csr_matvec, 2 → _csr_matmat; anything else (3+ dims, or 0-d) falls through to NotImplementedError with the offending shape. Batch the product or use BCOO/BCSR which support batched operands.

Source

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

  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=}")
    obj.__dict__.update(**aux_data)
    return obj


@tree_util.register_pytree_node_class
class CSC(JAXSparse):
  """Experimental CSC matrix implemented in JAX; API subject to change."""
  data: jax.Array

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.vmap over the leading batch dims: jax.vmap(lambda w: csr @ w)(tensor)
2. Reshape/squeeze the operand to ndim 1 or 2
3. Switch to BCOO (bcoo_matmul) or BCSR for natively batched sparse-dense matmul

Example fix

# before
c = csr_mat @ W  # W.shape == (4, 8, 8) -> NotImplementedError

# after
c = jax.vmap(lambda w: csr_mat @ w)(W)

assert jnp.asarray(other).ndim in (1, 2), f'ndim={other.ndim} unsupported by CSR matmul'

def csr_matmul_shape_ok(other) -> bool:
    import jax.numpy as jnp
    return jnp.asarray(other).ndim in (1, 2)

try:
    c = csr_mat @ W
except NotImplementedError:
    c = jax.vmap(lambda w: csr_mat @ w)(W)

When it happens

Trigger: csr_mat @ tensor3d where tensor3d.ndim >= 3 (e.g. shape (B, N, M)), or '@' with a 0-d array.

Common situations: Passing a batched stack of matrices/vectors from a dense pipeline directly; leftover size-1 dims making a vector ndim 3; vectorized training loops that assume broadcasting like numpy matmul.

Related errors

• matmul with object of shape {other.shape}
• CSR must have ndim=2; got {shape=}
• matmul between two sparse objects.
• scaled_matmul requires all inputs to be 3-dimensional arrays
• scaled_matmul requires inputs a and b to have matching batch

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