jax-ml/jax · error · ValueError

Cannot map in_axis={axis} for BCSR array with n_batch={val.n

Error message

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.

What it means

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.

Source

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

    if n_dense != 0 or n_batch != 0:
      raise NotImplementedError("BCSR from_scipy_sparse with nonzero n_dense/n_batch.")

    if mat.ndim != 2:
      raise ValueError(f"BCSR from_scipy_sparse requires 2D array; {mat.ndim}D is given.")

    mat = mat.tocsr()
    data = jnp.asarray(mat.data)
    indices = jnp.asarray(mat.indices).astype(index_dtype or jnp.int32)
    indptr = jnp.asarray(mat.indptr).astype(index_dtype or jnp.int32)
    return cls((data, indices, indptr), shape=mat.shape)

#--------------------------------------------------------------------
# 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)),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure the BCSR actually has batch dims: construct the source BCOO/BCSR with n_batch=1 (or use bcoo.reshape_to_batched) before vmap
  2. Map over an existing batch axis index < n_batch instead of a sparse/dense axis
  3. Reexpress the operation with sparse.bcsr_* batched primitives (e.g. bcsr_matmul) rather than vmap
  4. Fall back to BCOO, whose vmap support is broader

Example fix

# before
x = BCSR.from_bcoo(bcoo.bcoo_fromdense(x3d))  # n_batch=0
f = jax.vmap(lambda m, v: m @ v, in_axes=(0, None))
f(x, v)  # ValueError: Cannot map in_axis=0

# after
x = BCSR.from_bcoo(bcoo.reshape_to_batched(bcoo.bcoo_fromdense(x3d), 1))
f(x, v)
Defensive patterns

Strategy: validation

Validate before calling

assert isinstance(in_axes, int) is False or bcsr_arr.n_batch > in_axes, \
    f'in_axis must be < n_batch={bcsr_arr.n_batch}'

Type guard

def bcsr_vmap_axis_ok(arr, axis) -> bool:
    return axis is None or (0 <= axis < arr.n_batch)

Try / catch

try:
    f = jax.vmap(fn, in_axes=0)(x)
except ValueError:
    x = BCSR.from_bcoo(bcoo.reshape_to_batched(bcoo.BCOO.from_bcsr... , 1))
    f = jax.vmap(fn, in_axes=0)(x)

Prevention

When it happens

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

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

Related errors

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