jax-ml/jax · error · ValueError

data batch dimensions not compatible for {data.shape=}, {sha

Error message

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

What it means

For a valid BCOO, each batch dimension of data must equal the corresponding matrix batch dim or be 1 (broadcastable). If data.shape's leading batch dims don't match shape's batch dims, _validate_bcoo raises ValueError.

Source

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

class BCOOProperties(NamedTuple):
  n_batch: int
  n_sparse: int
  n_dense: int
  nse: int

class Buffer(Protocol):
  @property
  def shape(self) -> Shape: ...
  @property
  def dtype(self) -> Any: ...


def _validate_bcoo(data: Buffer, indices: Buffer, shape: Sequence[int]) -> BCOOProperties:
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Construct via high-level APIs: BCOO.fromdense, sparsify, or jax.sparse.bcoo_fromdense
Ensure data.shape[:n_batch] matches shape[:n_batch] elementwise or is 1
Use bcoo_sum_duplicates / bcoo_eliminate_zeros and reshape helpers instead of manual buffer surgery

Example fix

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

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

Build BCOOs via fromdense/sparsify, not raw buffers
After slicing, re-validate data batch dims against shape

When it happens

Trigger: Constructing a BCOO directly with BCOO((data, indices)) where data has shape like (3, nse, ...) but the declared shape has batch dim 2, or batch dims of data not in {1, batch_size}.

Common situations: Manually building BCOO buffers instead of using fromdense/sparsify; reshaping or slicing data without matching indices; mismatched batch sizes between data and indices.

Related errors

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