{"record":{"id":"941749260498815b","repo":"jax-ml/jax","slug":"batch-dims-out-of-range-for-indices-with-n-bat","errorCode":null,"errorMessage":"{batch_dims=} out of range for indices with {n_batch=}","messagePattern":"(.+?) out of range for indices with (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcoo.py","lineNumber":491,"sourceCode":"def _bcoo_extract_batching_rule(batched_args, batch_dims, *, assume_unique):\n  indices, arr = batched_args\n  assert any(b is not None for b in batch_dims)\n  if batch_dims[0] is None:\n    bdim = batch_dims[1]\n    indices = lax.expand_dims(indices, (bdim,))\n  elif batch_dims[1] is None:\n    # TODO(jakevdp) can we handle this case without explicit broadcasting?\n    bdim = batch_dims[0]\n    result_shape = list(arr.shape)\n    result_shape.insert(bdim, indices.shape[bdim])\n    arr = lax.broadcast_in_dim(arr, result_shape, (bdim,))\n  else:\n    if batch_dims[0] != batch_dims[1]:\n      raise NotImplementedError(\"bcoo_extract with unequal batch dimensions.\")\n    bdim = batch_dims[0]\n  n_batch = indices.ndim - 2\n  if bdim >= n_batch:\n    raise ValueError(f\"{batch_dims=} out of range for indices with {n_batch=}\")\n  return _bcoo_extract(indices, arr, assume_unique=assume_unique), bdim\n\nad.defjvp(bcoo_extract_p, None, _bcoo_extract_jvp)\nad.primitive_transposes[bcoo_extract_p] = _bcoo_extract_transpose\nbatching.primitive_batchers[bcoo_extract_p] = _bcoo_extract_batching_rule\nmlir.register_lowering(bcoo_extract_p, mlir.lower_fun(\n    _bcoo_extract_impl, multiple_results=False))\n\n#----------------------------------------------------------------------\n# bcoo_transpose\n# transpose of a BCOO array\n\nbcoo_transpose_p = core.Primitive('bcoo_transpose')\nbcoo_transpose_p.multiple_results = True\n\ndef bcoo_transpose(mat: BCOO, *, permutation: Sequence[int]) -> BCOO:\n  \"\"\"Transpose a BCOO-format array.\n","sourceCodeStart":473,"sourceCodeEnd":509,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcoo.py#L473-L509","documentation":"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.","triggerScenarios":"vmap-ing over the nse axis (the per-element axis of the sparse representation) of indices, producing bdim >= n_batch in the batching rule.","commonSituations":"Mapping a function over each stored element of a sparse matrix rather than over batch dims; treating indices' rows as a batch.","solutions":["Map only over batch dimensions; per-element operations should use the data buffer directly (e.g. vmap over sparr.data)","Restructure so the mapped axis is a leading batch dim of the BCOO"],"exampleFix":"// before\nf = jax.vmap(lambda idx: _bcoo_extract(idx, dense))(bcoo.indices)  # maps nse axis\n// after\nf = jax.vmap(lambda d: d)(bcoo.data)  # operate on stored values directly","handlingStrategy":"validation","validationCode":"n_batch = bcoo.indices.ndim - 2\nassert all(b is None or b < n_batch for b in in_axes), 'mapped axis is not a batch dim'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Map over batch dims only; operate on .data for per-element work"],"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"}