{"record":{"id":"97f07c6416934aaa","repo":"jax-ml/jax","slug":"bcsr-cannot-add-out-axis-axis-for-bcsr-array-wi","errorCode":null,"errorMessage":"BCSR: cannot add out_axis={axis} for BCSR array with n_batch={elt.n_batch}. BCSR batch axes must be a contiguous block of leading dimensions.","messagePattern":"BCSR: cannot add out_axis=(.+?) for BCSR array with n_batch=(.+?)\\. BCSR batch axes must be a contiguous block of leading dimensions\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/sparse/bcsr.py","lineNumber":1022,"sourceCode":"# 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)),\n              shape=elt.shape[:axis] + (axis_size,) + elt.shape[axis:])\n\nbatching.register_vmappable(BCSR, int, int, _bcsr_to_elt, _bcsr_from_elt, None)\n","sourceCodeStart":1004,"sourceCodeEnd":1031,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/sparse/bcsr.py#L1004-L1031","documentation":"When jax.vmap produces a BCSR output, the new mapped axis must be inserted at or before the existing batch dimensions (axis <= n_batch) so batch dims stay a contiguous leading block — a structural requirement of the BCSR buffer layout. _bcsr_from_elt raises ValueError when out_axis would place the new axis after the batch block.","triggerScenarios":"jax.vmap with out_axes=k where k > bcsr_output.n_batch (e.g. out_axes=-1 on a BCSR with n_batch=0, or out_axes=2 on n_batch=1 output). The vmapped function returns a BCSR and the caller wants the mapped dim in a non-leading position.","commonSituations":"Composing vmap twice with mismatched in_axes/out_axes; requesting out_axes=-1 out of dense-pipeline habit; stacking sparse results along the last axis.","solutions":["Use out_axes in [0, n_batch] (typically 0) so the new axis lands in the batch block","If a trailing mapped axis is required, vmap with out_axes=0 and then transpose the DENSE result after todense(), or use BCOO which is more flexible","Reorder axes after converting to dense instead of forcing out_axis on the sparse object"],"exampleFix":"# before\nf = jax.vmap(make_bcsr, in_axes=0, out_axes=-1)  # ValueError: cannot add out_axis\n\n# after\nf = jax.vmap(make_bcsr, in_axes=0, out_axes=0)","handlingStrategy":"validation","validationCode":"# ensure out_axes <= n_batch of the function's BCSR output\nout_axes = 0  # always valid","typeGuard":null,"tryCatchPattern":"try:\n    stacked = jax.vmap(fn, out_axes=-1)(x)\nexcept ValueError:\n    stacked = jax.vmap(fn, out_axes=0)(x)","preventionTips":["Use out_axes=0 for vmapped functions returning BCSR","Transpose after todense if a trailing axis is required","Remember BCSR batch dims must be leading and contiguous"],"tags":["jax","sparse","bcsr","vmap","out-axes"],"backgroundTag":"vmap-axis-not-batched","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}