{"record":{"id":"8a6e4de100eb5e4a","repo":"jax-ml/jax","slug":"cannot-map-in-axis-axis-for-bcsr-array-with-n-ba","errorCode":null,"errorMessage":"Cannot map in_axis={axis} for BCSR array with n_batch={val.n_batch}. in_axes for batched BCSR operations must correspond to a batched dimension.","messagePattern":"Cannot map in_axis=(.+?) for BCSR array with n_batch=(.+?)\\. in_axes for batched BCSR operations must correspond to a batched dimension\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcsr.py","lineNumber":1009,"sourceCode":"    if n_dense != 0 or n_batch != 0:\n      raise NotImplementedError(\"BCSR from_scipy_sparse with nonzero n_dense/n_batch.\")\n\n    if mat.ndim != 2:\n      raise ValueError(f\"BCSR from_scipy_sparse requires 2D array; {mat.ndim}D is given.\")\n\n    mat = mat.tocsr()\n    data = jnp.asarray(mat.data)\n    indices = jnp.asarray(mat.indices).astype(index_dtype or jnp.int32)\n    indptr = jnp.asarray(mat.indptr).astype(index_dtype or jnp.int32)\n    return cls((data, indices, indptr), shape=mat.shape)\n\n#--------------------------------------------------------------------\n# vmappable handlers\ndef _bcsr_to_elt(cont, _, val, axis):\n  if axis is None:\n    return val\n  if axis >= val.n_batch:\n    raise ValueError(f\"Cannot map in_axis={axis} for BCSR array with n_batch=\"\n                     f\"{val.n_batch}. in_axes for batched BCSR operations must \"\n                     \"correspond to a batched dimension.\")\n  return BCSR((cont(val.data, axis),\n               cont(val.indices, axis),\n               cont(val.indptr, axis)),\n              shape=val.shape[:axis] + val.shape[axis + 1:])\n\n\ndef _bcsr_from_elt(cont, axis_size, elt, axis):\n  if axis is None:\n    return elt\n  if axis > elt.n_batch:\n    raise ValueError(f\"BCSR: cannot add out_axis={axis} for BCSR array with \"\n                     f\"n_batch={elt.n_batch}. BCSR batch axes must be a \"\n                     \"contiguous block of leading dimensions.\")\n  return BCSR((cont(axis_size, elt.data, axis),\n               cont(axis_size, elt.indices, axis),\n               cont(axis_size, elt.indptr, axis)),","sourceCodeStart":991,"sourceCodeEnd":1027,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcsr.py#L991-L1027","documentation":"jax.vmap over a BCSR array can only map along batch dimensions, because the sparse data/indices/indptr buffers only have an explicit leading-batch layout. _bcsr_to_elt (the vmappable handler) raises ValueError when the requested in_axis is >= n_batch, i.e. it points into the sparse or dense dimensions.","triggerScenarios":"jax.vmap(fn, in_axes=k)(bcsr_array) where k >= bcsr_array.n_batch (n_batch is often 0, so any integer in_axis fails). Also mapping a per-sample function over a BCSR whose batch dims were never declared.","commonSituations":"Building a BCSR from a 3D BCOO without n_batch, then vmap-ing over axis 0; assuming vmap over sparse dims works like dense arrays; migrating dense vmap pipelines to sparse.","solutions":["Ensure the BCSR actually has batch dims: construct the source BCOO/BCSR with n_batch=1 (or use bcoo.reshape_to_batched) before vmap","Map over an existing batch axis index < n_batch instead of a sparse/dense axis","Reexpress the operation with sparse.bcsr_* batched primitives (e.g. bcsr_matmul) rather than vmap","Fall back to BCOO, whose vmap support is broader"],"exampleFix":"# before\nx = BCSR.from_bcoo(bcoo.bcoo_fromdense(x3d))  # n_batch=0\nf = jax.vmap(lambda m, v: m @ v, in_axes=(0, None))\nf(x, v)  # ValueError: Cannot map in_axis=0\n\n# after\nx = BCSR.from_bcoo(bcoo.reshape_to_batched(bcoo.bcoo_fromdense(x3d), 1))\nf(x, v)","handlingStrategy":"validation","validationCode":"assert isinstance(in_axes, int) is False or bcsr_arr.n_batch > in_axes, \\\n    f'in_axis must be < n_batch={bcsr_arr.n_batch}'","typeGuard":"def bcsr_vmap_axis_ok(arr, axis) -> bool:\n    return axis is None or (0 <= axis < arr.n_batch)","tryCatchPattern":"try:\n    f = jax.vmap(fn, in_axes=0)(x)\nexcept ValueError:\n    x = BCSR.from_bcoo(bcoo.reshape_to_batched(bcoo.BCOO.from_bcsr... , 1))\n    f = jax.vmap(fn, in_axes=0)(x)","preventionTips":["Always declare batch dims (n_batch) on BCSR that will be vmapped","Prefer sparse batched primitives (bcsr_matmul) over vmap","Check arr.n_batch > 0 before vmap with integer in_axes"],"tags":["jax","sparse","bcsr","vmap","batching"],"backgroundTag":"vmap-axis-not-batched","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}