jax-ml/jax · error · NotImplementedError

matmul with object of shape {other.shape}

Error message

matmul with object of shape {other.shape}

What it means

COO.__matmul__ only accepts dense right operands with ndim 1 (matvec) or 2 (matmat). A dense operand with 3+ dimensions (or an otherwise odd ndim) hits the final NotImplementedError. Reduce/reshape the operand, or use BCOO's batched matmul.

Source

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

    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.

  Args:
    mat : COO matrix
  Returns:
    mat_dense: dense version of ``mat``
  """
  return _coo_todense(mat.data, mat.row, mat.col, spinfo=mat._info)

def _coo_todense(data: Array, row: Array, col: Array, *, spinfo: COOInfo) -> Array:
  """Convert CSR-format sparse matrix to a dense matrix.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Squeeze/reshape the operand to ndim 1 or 2 (e.g. other.reshape(other.shape[-2:]) per batch or other.squeeze())
2. Batch the product with jax.vmap over leading dims
3. Use BCOO and bcoo.bcoo_matmul / sparse.sparsify for natively batched sparse-dense products

Example fix

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

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

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

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

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

When it happens

Trigger: coo_mat @ dense_3d_tensor, or '@' with an operand whose ndim is not 1 or 2 (e.g. a batched stack of vectors with shape (B, N, 1)).

Common situations: Feeding a batched weight tensor directly instead of looping/vmapping; leftover extra size-1 dims from broadcasting in a dense pipeline (shape (N, 1) vs (N,)).

Related errors

• Unsupported shape: {shape}
• COO must have ndim=2; got {shape=}
• matmul between two sparse objects.
• matmul with object of shape {other.shape}
• scaled_matmul requires all inputs to be 3-dimensional arrays

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