jax-ml/jax · error · ValueError

indices batch dimensions not compatible for {indices.shape=}

Error message

indices batch dimensions not compatible for {indices.shape=}, {shape=}

What it means

Each batch dim of the indices array of a BCOO must equal the corresponding matrix batch dim or be 1 (broadcastable). Otherwise _validate_bcoo_indices raises ValueError.

Source

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

  props = _validate_bcoo_indices(indices, shape)
  n_batch, n_sparse, n_dense, nse = props
  shape = tuple(shape)
  if any(s1 not in (1, s2) for s1, s2 in safe_zip(data.shape[:n_batch], shape[:n_batch])):
    raise ValueError(f"data batch dimensions not compatible for {data.shape=}, {shape=}")
  if data.shape[n_batch:] != (nse,) + shape[n_batch + n_sparse:]:
    raise ValueError(f"Invalid {data.shape=} for {nse=}, {n_batch=}, {n_dense=}")
  return props


def _validate_bcoo_indices(indices: Buffer, shape: Sequence[int]) -> BCOOProperties:
  assert jnp.issubdtype(indices.dtype, jnp.integer)
  shape = tuple(shape)
  nse, n_sparse = indices.shape[-2:]
  n_batch = len(indices.shape) - 2
  n_dense = len(shape) - n_batch - n_sparse
  assert n_dense >= 0
  if any(s1 not in (1, s2) for s1, s2 in safe_zip(indices.shape[:n_batch], shape[:n_batch])):
    raise ValueError(f"indices batch dimensions not compatible for {indices.shape=}, {shape=}")
  if indices.shape[n_batch:] != (nse, n_sparse):
    raise ValueError(f"Invalid ={indices.shape=} for {nse=}, {n_batch=}, {n_dense=}")
  return BCOOProperties(n_batch=n_batch, n_sparse=n_sparse, n_dense=n_dense, nse=nse)


#----------------------------------------------------------------------
# bcoo_todense

bcoo_todense_p = core.Primitive('bcoo_todense')

def bcoo_todense(mat: BCOO) -> Array:
  """Convert batched sparse matrix to a dense matrix.

  Args:
    mat: BCOO matrix.

  Returns:
    mat_dense: dense version of ``mat``.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Align indices batch dims with shape batch dims (or 1) before constructing BCOO
Use bcoo_concatenate / BCOO methods rather than np.concatenate on buffers

Example fix

// before
m = BCOO((data, indices), shape=(3, 5, 5))  # indices.shape[0]==4
// after
m = BCOO((data, indices), shape=(4, 5, 5))

Defensive patterns

Strategy: validation

Validate before calling

n_batch = indices.ndim - 2
assert all(d in (1, s) for d, s in zip(indices.shape[:n_batch], shape[:n_batch]))

Prevention

Use bcoo_concatenate for combining sparse arrays
Check indices batch dims after any manual slicing

When it happens

Trigger: Directly constructing or slicing BCOO indices whose leading (batch) dims mismatch the declared shape's batch dims, e.g. indices.shape=(4,n,d) with shape batch dim 3.

Common situations: Manual indices manipulation; concatenating BCOOs with different batch sizes; off-by-one batch slicing.

Related errors

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