jax-ml/jax · error · ValueError

Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}

Error message

Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}

What it means

When vmap-ing bcoo_fromdense, the batched axis of the dense input must be within [0, n_batch]. The code even has an off-by-one message bug (checks 0 <= bdim <= n_batch but the message says 0 < bdim). Raised for out-of-range batch dims.

Source

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

  return primals_out, tangents_out

def _bcoo_fromdense_transpose(ct, M, *, nse, n_batch, n_dense, index_dtype):
  data, indices = ct
  n_sparse = M.ndim - n_batch - n_dense
  assert data.shape == M.shape[:n_batch] + (nse,) + M.shape[n_batch + n_sparse:]
  assert indices.shape == M.shape[:n_batch] + (n_sparse, nse)
  assert indices.dtype == index_dtype
  if isinstance(indices, ad.Zero):
    raise ValueError("Cannot transpose with respect to sparse indices")
  assert ad.is_undefined_primal(M)
  return _bcoo_todense(data, indices, spinfo=SparseInfo(M.aval.shape))

def _bcoo_fromdense_batching_rule(batched_args, batch_dims, *, nse, n_batch, n_dense, index_dtype):
  M, = batched_args
  bdim, = batch_dims
  if not (0 <= bdim <= n_batch):
    raise ValueError(f"Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}")
  return _bcoo_fromdense(M, nse=nse, n_batch=n_batch + 1, n_dense=n_dense, index_dtype=index_dtype), (bdim, bdim)

ad.primitive_jvps[bcoo_fromdense_p] = _bcoo_fromdense_jvp
ad.primitive_transposes[bcoo_fromdense_p] = _bcoo_fromdense_transpose
batching.primitive_batchers[bcoo_fromdense_p] = _bcoo_fromdense_batching_rule
mlir.register_lowering(bcoo_fromdense_p, mlir.lower_fun(
    _bcoo_fromdense_impl, multiple_results=True))

#----------------------------------------------------------------------
# bcoo_extract

bcoo_extract_p = core.Primitive('bcoo_extract')


def bcoo_extract(sparr: BCOO, arr: ArrayLike, *, assume_unique: bool | None = None) -> BCOO:
  """Extract values from a dense array according to the sparse array's indices.

  Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Declare the mapped dimension as a batch dim: pass n_batch including it in BCOO.fromdense / bcoo_fromdense
2. Move the mapped axis to position 0 and set n_batch accordingly
3. Do the dense-to-sparse conversion outside vmap, then vmap over the resulting BCOO

Example fix

// before
f = jax.vmap(lambda x: BCOO.fromdense(x, nse=4))(batched_dense)  # n_batch mismatch
// after
f = BCOO.fromdense(batched_dense, nse=4, n_batch=1)  # batch conversion directly

assert 0 <= mapped_axis <= n_batch_declared, 'bdim outside batch dims'

When it happens

Trigger: jax.vmap over an axis beyond the declared n_batch when converting dense to BCOO, e.g. vmap-ing a fromdense call where the mapped dim falls in sparse dimensions.

Common situations: vmap over rows of a matrix being converted elementwise to sparse; converting dense arrays whose leading dims aren't declared as batch dims.

Related errors

• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b
• bcoo_extract with unequal batch dimensions.
• {batch_dims=} out of range for indices with {n_batch=}
• 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/ff299ade2dafd953. Report an issue: GitHub.