jax-ml/jax · error · NotImplementedError

bcoo_extract with unequal batch dimensions.

Error message

bcoo_extract with unequal batch dimensions.

What it means

When vmap-ing bcoo_extract with both indices and dense array batched, their batch dims must be equal; unequal batch dims (e.g. one batched at axis 0, the other at axis 1) are not implemented and raise NotImplementedError.

Source

Thrown at jax/experimental/sparse/bcoo.py:487

    raise ValueError("Cannot transpose with respect to sparse indices")
  assert ct.dtype == arr.aval.dtype
  return indices, _bcoo_todense(ct, indices, spinfo=SparseInfo(arr.aval.shape))

def _bcoo_extract_batching_rule(batched_args, batch_dims, *, assume_unique):
  indices, arr = batched_args
  assert any(b is not None for b in batch_dims)
  if batch_dims[0] is None:
    bdim = batch_dims[1]
    indices = lax.expand_dims(indices, (bdim,))
  elif batch_dims[1] is None:
    # TODO(jakevdp) can we handle this case without explicit broadcasting?
    bdim = batch_dims[0]
    result_shape = list(arr.shape)
    result_shape.insert(bdim, indices.shape[bdim])
    arr = lax.broadcast_in_dim(arr, result_shape, (bdim,))
  else:
    if batch_dims[0] != batch_dims[1]:
      raise NotImplementedError("bcoo_extract with unequal batch dimensions.")
    bdim = batch_dims[0]
  n_batch = indices.ndim - 2
  if bdim >= n_batch:
    raise ValueError(f"{batch_dims=} out of range for indices with {n_batch=}")
  return _bcoo_extract(indices, arr, assume_unique=assume_unique), bdim

ad.defjvp(bcoo_extract_p, None, _bcoo_extract_jvp)
ad.primitive_transposes[bcoo_extract_p] = _bcoo_extract_transpose
batching.primitive_batchers[bcoo_extract_p] = _bcoo_extract_batching_rule
mlir.register_lowering(bcoo_extract_p, mlir.lower_fun(
    _bcoo_extract_impl, multiple_results=False))

#----------------------------------------------------------------------
# bcoo_transpose
# transpose of a BCOO array

bcoo_transpose_p = core.Primitive('bcoo_transpose')
bcoo_transpose_p.multiple_results = True

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use matching in_axes for both operands, e.g. in_axes=0
2. Move both operands' mapped axes to the same position with jnp.moveaxis before vmap

Example fix

// before
out = jax.vmap(bcoo_extract, in_axes=(0, 1))(bcoo_batch, dense_batch)
// after
dense_batch = jnp.moveaxis(dense_batch, 1, 0)
out = jax.vmap(bcoo_extract, in_axes=0)(bcoo_batch, dense_batch)

assert in_axes[0] == in_axes[1] or None in in_axes, 'batch dims must match'

When it happens

Trigger: jax.vmap(bcoo_extract, in_axes=(0, 1)) or any vmap configuration yielding different batch_dims for indices and arr inside the batching rule.

Common situations: Batching the sparse operand along one axis and the dense operand along another; in_axes mismatch when mapping lists of matrices.

Related errors

• batch_dims must be None or satisfy 0 < dim < n_batch. Got {b
• {batch_dims=} out of range for indices with {n_batch=}
• Expected 0 < bdim <= n_batch; got {bdim=}, {n_batch=}
• transpose permutation cannot permute batch axes with non-bat
• Cannot map in_axis={axis} for BCSR array with n_batch={val.n

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