jax-ml/jax · error · ValueError

indices batch dimensions not compatible for

Error message

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

What it means

BCOO validation checks that the indices buffer's leading (batch) dimensions are compatible with the declared dense shape — each batch extent must be 1 or equal to the corresponding shape entry. It fires when broadcasting the batch dims cannot produce the declared shape.

Solutions

1. Align indices batch dims with shape batch dims (or 1) before constructing BCOO
2. 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))

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

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

• data batch dimensions not compatible for
• Invalid for
• Invalid = for
• shape mismatch
• batch_dims must be None or satisfy 0 < dim < n_batch. Got

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

Appendix: 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)