{"record":{"id":"107deaebba49c05c","repo":"jax-ml/jax","slug":"batch-dims-must-be-none-or-satisfy-0-dim-n-bat","errorCode":null,"errorMessage":"batch_dims must be None or satisfy 0 < dim < n_batch. Got {batch_dims=} for {n_batch=}.","messagePattern":"batch_dims must be None or satisfy 0 < dim < n_batch\\. Got (.+?) for (.+?)\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":67,"sourceCode":"from jax._src.lax.lax import (\n  _const, _unbroadcast, ranges_like, remaining, _dot_general_batch_dim_nums,\n  DotDimensionNumbers)\nfrom jax._src.lax.slicing import GatherDimensionNumbers, GatherScatterMode\nfrom jax._src.numpy.setops import _unique\nfrom jax._src.typing import Array, ArrayLike, DTypeLike\nfrom jax._src.util import canonicalize_axis\n\n\nCUSPARSE_DATA_DTYPES = [np.float32, np.float64, np.complex64, np.complex128]\nCUSPARSE_INDEX_DTYPES = [np.int32]\n\n\ndef _bcoo_batch_dims_to_front(batched_args, batch_dims, spinfo, batch_size=None):\n  data, indices = batched_args\n  data_bdim, indices_bdim = batch_dims\n  n_batch = indices.ndim - 2 + bool(indices_bdim is None)\n  if not all(b is None or 0 <= b < n_batch for b in batch_dims):\n    raise NotImplementedError(\"batch_dims must be None or satisfy 0 < dim < n_batch. \"\n                              f\"Got {batch_dims=} for {n_batch=}.\")\n  batched_data, batched_indices = (\n      lax.expand_dims(arg, [0]) if bdim is None else jnp.moveaxis(arg, bdim, 0)\n      for arg, bdim in [(data, data_bdim), (indices, indices_bdim)])\n  if batch_size is None:\n    batch_size = max(arg.shape[dim] for arg, dim in zip((data, indices), batch_dims) if dim is not None)\n  batched_spinfo = SparseInfo((batch_size, *spinfo.shape),\n                              indices_sorted=spinfo.indices_sorted,\n                              unique_indices=spinfo.unique_indices)\n  return batched_data, batched_indices, batched_spinfo\n\n\n#----------------------------------------------------------------------\n# BCOO primitives: batched extension of COO.\n\ndef _bcoo_set_nse(mat: BCOO, nse: int) -> BCOO:\n  \"\"\"Return a copy of `mat` with the specified nse.\n  Note that if nse < mat.nse, this will potentially discard data.","sourceCodeStart":49,"sourceCodeEnd":85,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L49-L85","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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)","Move the mapped axis to the front batch dimension before vmap","Use BCOO with nse broadcast/adjustment if you need per-batch sparsity"],"exampleFix":"// before\nf = jax.vmap(lambda m: m.todense())(bcoo_with_nse_axis_batched)\n// after\nm = bcoo.reshape(...)  # ensure batch dim is axis 0\nf = jax.vmap(lambda x: x.todense())(m)","handlingStrategy":"validation","validationCode":"n_batch = bcoo.indices.ndim - 2\nassert all(b is None or 0 <= b < n_batch for b in in_axes_tuple), 'bad batch dims'","typeGuard":null,"tryCatchPattern":"try:\n    jax.vmap(f, in_axes=...)(bcoo)\nexcept NotImplementedError as e:\n    # restructure so mapped axis is a batch dim\n    raise","preventionTips":["Only vmap over BCOO batch dimensions","Construct batched BCOOs with the mapped axis as a leading batch dim"],"tags":["jax","sparse","bcoo","vmap","batching"],"backgroundTag":"vmap-batch-dim-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}