jax-ml/jax · error · ValueError

BCSR: cannot add out_axis={axis} for BCSR array with n_batch

Error message

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.

What it means

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.

Source

Thrown at jax/experimental/sparse/bcsr.py:1022

# vmappable handlers
def _bcsr_to_elt(cont, _, val, axis):
  if axis is None:
    return val
  if axis >= val.n_batch:
    raise ValueError(f"Cannot map in_axis={axis} for BCSR array with n_batch="
                     f"{val.n_batch}. in_axes for batched BCSR operations must "
                     "correspond to a batched dimension.")
  return BCSR((cont(val.data, axis),
               cont(val.indices, axis),
               cont(val.indptr, axis)),
              shape=val.shape[:axis] + val.shape[axis + 1:])


def _bcsr_from_elt(cont, axis_size, elt, axis):
  if axis is None:
    return elt
  if axis > elt.n_batch:
    raise ValueError(f"BCSR: cannot add out_axis={axis} for BCSR array with "
                     f"n_batch={elt.n_batch}. BCSR batch axes must be a "
                     "contiguous block of leading dimensions.")
  return BCSR((cont(axis_size, elt.data, axis),
               cont(axis_size, elt.indices, axis),
               cont(axis_size, elt.indptr, axis)),
              shape=elt.shape[:axis] + (axis_size,) + elt.shape[axis:])

batching.register_vmappable(BCSR, int, int, _bcsr_to_elt, _bcsr_from_elt, None)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use out_axes in [0, n_batch] (typically 0) so the new axis lands in the batch block
2. 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
3. Reorder axes after converting to dense instead of forcing out_axis on the sparse object

Example fix

# before
f = jax.vmap(make_bcsr, in_axes=0, out_axes=-1)  # ValueError: cannot add out_axis

# after
f = jax.vmap(make_bcsr, in_axes=0, out_axes=0)

# ensure out_axes <= n_batch of the function's BCSR output
out_axes = 0  # always valid

try:
    stacked = jax.vmap(fn, out_axes=-1)(x)
except ValueError:
    stacked = jax.vmap(fn, out_axes=0)(x)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Cannot map in_axis={axis} for BCSR array with n_batch={val.n
• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b
• Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}
• bcoo_extract with unequal batch dimensions.
• {batch_dims=} out of range for indices with {n_batch=}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/97f07c6416934aaa. Report an issue: GitHub.