jax-ml/jax · error · NotImplementedError

bcoo_multiply_sparse: arrays with differing numbers of dense

Error message

bcoo_multiply_sparse: arrays with differing numbers of dense dimensions: {lhs}, {rhs}

What it means

bcoo_multiply_sparse supports differing numbers of batch dimensions (it takes the min via vmap) but not differing numbers of dense dimensions; if lhs.n_dense != rhs.n_dense it raises NotImplementedError because no layout conversion is performed.

Source

Thrown at jax/experimental/sparse/bcoo.py:2219

  """
  out_data, out_indices, out_shape = _bcoo_multiply_sparse(
      lhs.data, lhs.indices, rhs.data, rhs.indices, lhs_spinfo=lhs._info,
      rhs_spinfo=rhs._info)
  return BCOO((out_data, out_indices), shape=out_shape)

def _bcoo_multiply_sparse(lhs_data: Array, lhs_indices: Array, rhs_data: Array, rhs_indices: Array, *,
                          lhs_spinfo: SparseInfo, rhs_spinfo: SparseInfo) -> tuple[Array, Array, Shape]:
  lhs_shape = lhs_spinfo.shape
  rhs_shape = rhs_spinfo.shape

  lhs = _validate_bcoo(lhs_data, lhs_indices, lhs_shape)
  rhs = _validate_bcoo(rhs_data, rhs_indices, rhs_shape)
  if len(lhs_shape) != len(rhs_shape):
    # Similar requirement as lax.mul:
    raise TypeError("bcoo_multiply_sparse: arrays must have same number of dimensions, "
                    f"got {lhs_shape}, {rhs_shape}")
  if lhs.n_dense != rhs.n_dense:
    raise NotImplementedError("bcoo_multiply_sparse: arrays with differing numbers of "
                              f"dense dimensions: {lhs}, {rhs}")
  n_batch = min(lhs.n_batch, rhs.n_batch)
  _mul = functools.partial(_bcoo_multiply_sparse_unbatched,
                           lhs_shape=lhs_shape[n_batch:],
                           rhs_shape=rhs_shape[n_batch:])
  _mul = nfold_vmap(_mul, n_batch)
  data, indices = _mul(lhs_data, lhs_indices, rhs_data, rhs_indices)
  return data, indices, jnp.broadcast_shapes(lhs_shape, rhs_shape)

def _bcoo_multiply_sparse_unbatched(lhs_data, lhs_indices, rhs_data, rhs_indices, *, lhs_shape, rhs_shape):
  lhs = _validate_bcoo(lhs_data, lhs_indices, lhs_shape)
  rhs = _validate_bcoo(rhs_data, rhs_indices, rhs_shape)
  assert (lhs.n_batch == 0) or (rhs.n_batch == 0)  # Ensured at call site above

  # TODO(jakevdp): this can be made more efficient by utilizing batch structure.
  if lhs.n_batch:
    lhs_data, lhs_indices = bcoo_update_layout(BCOO((lhs_data, lhs_indices), shape=lhs_shape), n_batch=0)._bufs
    lhs = _validate_bcoo(lhs_data, lhs_indices, lhs_shape)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Recreate one operand with matching dense dims: BCOO((data, indices), shape=...) with the same dense-dimension layout, or use mat.reshape to move dims between sparse/dense
2. Convert both operands from a canonical source (e.g. fromdense with the same dense_dimensions) so n_dense agrees
3. If semantically the shapes align, move the mismatched trailing axes so both have the same n_dense count

Example fix

# before
lhs = BCOO((d1, i1), shape=(m, n))          # n_dense=0
rhs = BCOO((d2, i2), shape=(m, n, k))       # n_dense=1
out = bcoo_multiply_sparse(...)  # NotImplementedError
# after
rhs2 = rhs.reshape(m, n * k) or rebuild lhs with n_dense=1 to match:
lhs = BCOO((d1[..., None], i1), shape=(m, n, 1))  # n_dense=1
out = lhs * rhs

if lhs.n_dense != rhs.n_dense:
    rhs = rhs.reshape(rhs.shape[:rhs.n_sparse], rhs.shape[rhs.n_sparse:])  # or rebuild with matching n_dense

from jax.experimental.sparse import BCOO
def compatible_layouts(a: BCOO, b: BCOO) -> bool:
    return a.n_dense == b.n_dense and a.ndim == b.ndim

try:
    out = lhs * rhs
except NotImplementedError as e:
    if 'dense dimensions' in str(e):
        rhs = BCOO.fromdense(rhs.todense(), n_batch=rhs.n_batch)  # realign layout
        out = lhs * rhs
    else:
        raise

When it happens

Trigger: Multiplying two BCOO arrays where one stores trailing dense dims (e.g. a block/ragged tensor with n_dense=1) and the other is a pure sparse matrix with n_dense=0.

Common situations: Mixing BCOO tensors created with different n_dense conventions (e.g. from _bcoo_fromdense with different dense_dimensions args); upgrading code that previously operated on uniform formats; multiplying a per-nonzero-vector tensor by a scalar-per-nonzero matrix.

Related errors

• Sparse metadata format not implemented for {operand_dtype=}
• Unsupported transforms: {a_sparse_metadata_transforms}
• Sparse meta layout loads unsupported.
• Sparse meta layout stores unsupported.
• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b

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