jax-ml/jax · error · TypeError

bcoo_multiply_sparse: arrays must have same number of dimens

Error message

bcoo_multiply_sparse: arrays must have same number of dimensions, got {lhs_shape}, {rhs_shape}

What it means

bcoo_multiply_sparse (element-wise sparse product used by BCOO multiplication) requires both operands to have the same number of dimensions, mirroring lax.mul's no-broadcasting-on-rank rule. A rank mismatch raises TypeError with both shapes.

Source

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

  Returns:
    An BCOO-format array containing the result.
  """
  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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape the lower-rank operand to match ranks, e.g. vec[None, :] or jnp.reshape before building BCOO
2. For matrix-vector products use bcoo_matvec / mat.dot(vec) instead of element-wise multiply
3. Rebuild both BCOO operands with consistent ndim and shape metadata

Example fix

# before
result = bcoo_multiply_sparse(mat_2d, vec_1d.data, vec_1d.indices, ..., shapes)
# after
vec = sparse.BCOO((v_data, v_indices), shape=(1, n))  # match 2-D rank
result = mat * vec  # or reshape appropriately

if lhs.ndim != rhs.ndim:
    raise ValueError(f'rank mismatch: {lhs.ndim} vs {rhs.ndim}; reshape first')

from jax.experimental.sparse import BCOO
def same_rank(a: BCOO, b: BCOO) -> bool:
    return len(a.spinfo.shape) == len(b.spinfo.shape)

When it happens

Trigger: Calling bcoo_multiply_sparse (or BCOO * BCOO / sparse multiply paths) where lhs is e.g. 2-D and rhs is 1-D, so len(lhs_shape) != len(rhs_shape).

Common situations: Multiplying a sparse matrix by a vector without reshaping; operands built from different pipelines with different rank; assuming NumPy broadcasting across ranks works here.

Related errors

• bcoo_slice: indices must have size mat.ndim={mat.ndim}
• bcoo_dynamic_slice: indices must have size mat.ndim={mat.ndi
• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b
• data batch dimensions not compatible for {data.shape=}, {sha
• Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}

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