jax-ml/jax · error · NotImplementedError

Unsupported shape: {shape}

Error message

Unsupported shape: {shape}

What it means

The COO sparse-matrix-times-dense GPU lowering only supports 2D matrices and 3D batched matrices. If the lhs shape has any other rank, the lowering cannot map it to a cuSPARSE routine and raises NotImplementedError.

Source

Thrown at jax/experimental/sparse/_lowerings.py:143

  return core.ShapedArray(
    shape=(shape[1] if transpose else shape[0], x.shape[1]),
    dtype=x.dtype)

def _coo_spmm_gpu_lowering(ctx, data, row, col, x, *, transpose, shape,
                           target_name_prefix):
  data_aval, row_aval, _, x_aval = ctx.avals_in
  nnz, = data_aval.shape
  _, Ccols = x_aval.shape

  batch_count = 1
  if len(shape) == 2:
    rows, cols = shape
  elif len(shape) == 3:
    batch_count, rows, cols = shape
    nnz = nnz // batch_count
  else:
    raise NotImplementedError(f"Unsupported shape: {shape}")

  # TODO(tianjianlu): use batch stride to trigger different mode of batch
  # computation. Currently batch_stride = 0 is not allowed because of the issue
  # in cusparse https://github.com/NVIDIA/CUDALibrarySamples/issues/81#issuecomment-1205562643
  # Set batch stride to be the matrix size for now.
  lhs_batch_stride = nnz
  B_rows = rows if transpose else cols
  rhs_batch_stride =  B_rows * Ccols

  buffer_size, opaque = _get_module(target_name_prefix).build_coo_matmat_descriptor(
      data_aval.dtype, x_aval.dtype, data_aval.dtype, row_aval.dtype,
      rows, cols, Ccols, nnz, transpose, batch_count, lhs_batch_stride,
      rhs_batch_stride)

  buffer_aval = core.ShapedArray(shape=(buffer_size,), dtype=np.int8)
  sub_ctx = ctx.replace(avals_out=[ctx.avals_out[0], buffer_aval])
  rule = ffi.ffi_lowering(f"{target_name_prefix}sparse_coo_matmat_ffi")
  return rule(sub_ctx, data, row, col, x, opaque=opaque)[:1]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to BCOO (sparsify with BCOO or arr.tobcoo()) which supports arbitrary batch/dense dims on GPU
2. Reshape so the sparse operand is exactly 2D (or 3D batched) before the matmul
3. Run on CPU, where the COO CPU lowering supports more shapes

Example fix

// before
out = coo_mat @ dense  # coo_mat has shape (B, N, M, K)
// after
bcoo_mat = coo_mat.reshape(...).tobcoo() if hasattr(coo_mat,'tobcoo') else coo_mat
out = bcoo_mat @ dense

shape = coo.shape
if len(shape) not in (2, 3) and jax.devices()[0].platform == 'gpu':
    coo = coo.reshape(...)  # or convert to BCOO

try:
    out = coo @ dense
except NotImplementedError:
    out = coo.tobcoo() @ dense

When it happens

Trigger: Calling matmul of a jax.experimental.sparse.COO with a dense array on GPU where the COO has rank other than 2 or 3 (e.g. extra dense dimensions making it 4D), via jnp.matmul / sp.dot dispatching to _coo_matmat_gpu_lowering.

Common situations: Using batched or higher-dimensional sparse arrays with COO format on GPU; COO lacks the dense-dimension and flexible batching support that BCOO has.

Related errors

• COO must have ndim=2; got {shape=}
• matmul with object of shape {other.shape}
• Sparse MMA not supported for M=64
• Only M=128 and M=64 are supported for MMA, but got M={m}
• Sparse MMA unsupported for f32

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