jax-ml/jax · error · ValueError

{batch_dims=} out of range for indices with {n_batch=}

Error message

{batch_dims=} out of range for indices with {n_batch=}

What it means

After resolving batching for bcoo_extract, the resulting batch dim must be one of the indices' batch dims (bdim < n_batch where n_batch = indices.ndim - 2). A bdim pointing into the nse or sparse-index axes is invalid.

Source

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

def _bcoo_extract_batching_rule(batched_args, batch_dims, *, assume_unique):
  indices, arr = batched_args
  assert any(b is not None for b in batch_dims)
  if batch_dims[0] is None:
    bdim = batch_dims[1]
    indices = lax.expand_dims(indices, (bdim,))
  elif batch_dims[1] is None:
    # TODO(jakevdp) can we handle this case without explicit broadcasting?
    bdim = batch_dims[0]
    result_shape = list(arr.shape)
    result_shape.insert(bdim, indices.shape[bdim])
    arr = lax.broadcast_in_dim(arr, result_shape, (bdim,))
  else:
    if batch_dims[0] != batch_dims[1]:
      raise NotImplementedError("bcoo_extract with unequal batch dimensions.")
    bdim = batch_dims[0]
  n_batch = indices.ndim - 2
  if bdim >= n_batch:
    raise ValueError(f"{batch_dims=} out of range for indices with {n_batch=}")
  return _bcoo_extract(indices, arr, assume_unique=assume_unique), bdim

ad.defjvp(bcoo_extract_p, None, _bcoo_extract_jvp)
ad.primitive_transposes[bcoo_extract_p] = _bcoo_extract_transpose
batching.primitive_batchers[bcoo_extract_p] = _bcoo_extract_batching_rule
mlir.register_lowering(bcoo_extract_p, mlir.lower_fun(
    _bcoo_extract_impl, multiple_results=False))

#----------------------------------------------------------------------
# bcoo_transpose
# transpose of a BCOO array

bcoo_transpose_p = core.Primitive('bcoo_transpose')
bcoo_transpose_p.multiple_results = True

def bcoo_transpose(mat: BCOO, *, permutation: Sequence[int]) -> BCOO:
  """Transpose a BCOO-format array.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Map only over batch dimensions; per-element operations should use the data buffer directly (e.g. vmap over sparr.data)
2. Restructure so the mapped axis is a leading batch dim of the BCOO

Example fix

// before
f = jax.vmap(lambda idx: _bcoo_extract(idx, dense))(bcoo.indices)  # maps nse axis
// after
f = jax.vmap(lambda d: d)(bcoo.data)  # operate on stored values directly

n_batch = bcoo.indices.ndim - 2
assert all(b is None or b < n_batch for b in in_axes), 'mapped axis is not a batch dim'

When it happens

Trigger: vmap-ing over the nse axis (the per-element axis of the sparse representation) of indices, producing bdim >= n_batch in the batching rule.

Common situations: Mapping a function over each stored element of a sparse matrix rather than over batch dims; treating indices' rows as a batch.

Related errors

• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b
• bcoo_extract with unequal batch dimensions.
• Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}
• transpose permutation cannot permute batch axes with non-bat
• Cannot map in_axis={axis} for BCSR array with n_batch={val.n

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