jax-ml/jax · error · NotImplementedError

batch_dims must be None or satisfy 0 < dim < n_batch. Got {b

Error message

batch_dims must be None or satisfy 0 < dim < n_batch. Got {batch_dims=} for {n_batch=}.

What it means

When batching (vmap) a BCOO primitive, the supplied batch_dims must each be None or a valid axis index within the batch dimensionality of the indices. Out-of-range batch dims raise NotImplementedError in _bcoo_batch_dims_to_front.

Source

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

from jax._src.lax.lax import (
  _const, _unbroadcast, ranges_like, remaining, _dot_general_batch_dim_nums,
  DotDimensionNumbers)
from jax._src.lax.slicing import GatherDimensionNumbers, GatherScatterMode
from jax._src.numpy.setops import _unique
from jax._src.typing import Array, ArrayLike, DTypeLike
from jax._src.util import canonicalize_axis


CUSPARSE_DATA_DTYPES = [np.float32, np.float64, np.complex64, np.complex128]
CUSPARSE_INDEX_DTYPES = [np.int32]


def _bcoo_batch_dims_to_front(batched_args, batch_dims, spinfo, batch_size=None):
  data, indices = batched_args
  data_bdim, indices_bdim = batch_dims
  n_batch = indices.ndim - 2 + bool(indices_bdim is None)
  if not all(b is None or 0 <= b < n_batch for b in batch_dims):
    raise NotImplementedError("batch_dims must be None or satisfy 0 < dim < n_batch. "
                              f"Got {batch_dims=} for {n_batch=}.")
  batched_data, batched_indices = (
      lax.expand_dims(arg, [0]) if bdim is None else jnp.moveaxis(arg, bdim, 0)
      for arg, bdim in [(data, data_bdim), (indices, indices_bdim)])
  if batch_size is None:
    batch_size = max(arg.shape[dim] for arg, dim in zip((data, indices), batch_dims) if dim is not None)
  batched_spinfo = SparseInfo((batch_size, *spinfo.shape),
                              indices_sorted=spinfo.indices_sorted,
                              unique_indices=spinfo.unique_indices)
  return batched_data, batched_indices, batched_spinfo


#----------------------------------------------------------------------
# BCOO primitives: batched extension of COO.

def _bcoo_set_nse(mat: BCOO, nse: int) -> BCOO:
  """Return a copy of `mat` with the specified nse.
  Note that if nse < mat.nse, this will potentially discard data.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. vmap only over an actual batch dimension of the BCOO; construct the BCOO with that axis as a batch dim (e.g. via BCOO.fromdense with n_batch, or reshape)
2. Move the mapped axis to the front batch dimension before vmap
3. Use BCOO with nse broadcast/adjustment if you need per-batch sparsity

Example fix

// before
f = jax.vmap(lambda m: m.todense())(bcoo_with_nse_axis_batched)
// after
m = bcoo.reshape(...)  # ensure batch dim is axis 0
f = jax.vmap(lambda x: x.todense())(m)

n_batch = bcoo.indices.ndim - 2
assert all(b is None or 0 <= b < n_batch for b in in_axes_tuple), 'bad batch dims'

try:
    jax.vmap(f, in_axes=...)(bcoo)
except NotImplementedError as e:
    # restructure so mapped axis is a batch dim
    raise

When it happens

Trigger: Using jax.vmap over BCOO ops (todense, transpose, dot_general, sort_indices, sum_duplicates) where the batched axis lies in the sparse/nse dimensions rather than the batch dimensions, or dims exceed n_batch.

Common situations: vmap-ing over the nse (number of stored elements) axis or a sparse axis, which JAX sparse does not support; spmd/vmap combinations producing unexpected batch_dims.

Related errors

• bcoo_extract with unequal batch dimensions.
• {batch_dims=} out of range for indices with {n_batch=}
• 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/107deaebba49c05c. Report an issue: GitHub.