jax-ml/jax · error · NotImplementedError

Currently only support batch_dim in [0, None], but got {dim=

Error message

Currently only support batch_dim in [0, None], but got {dim=}

What it means

Raised by _check_valid_batch_dims in the vmap batchers for cuDNN fused attention: only batch_dim 0 (batched) or None (unbatched) is supported per operand. vmapping over any other axis (e.g. heads or sequence) hits a batcher that cannot express that mapping for the fused cuDNN kernel.

Source

Thrown at jax/_src/cudnn/fused_attention_stablehlo.py:805

    backend_config=backend_config,
    operand_layouts=default_layouts(
      *[ir.RankedTensorType(operand.type).shape for operand in operands]),
    result_layouts=result_layouts,
  )
  dqkv = (hlo.transpose(out.results[0], grad_transpose_perm),
          hlo.transpose(out.results[1], grad_transpose_perm),
          hlo.transpose(out.results[2], grad_transpose_perm))
  # Only keep dQ, dK, dV and dBias here
  if has_dbias:
    return dqkv + (out.results[3],)
  else:
    return dqkv

# batcher
def _check_valid_batch_dims(bdims):
  for dim in bdims:
    if dim not in [0, None]:
      raise NotImplementedError(
        f"Currently only support batch_dim in [0, None], but got {dim=}")

def _broadcast_unbatched_args(batched_args, batch_dims, arg_idx):
  # Broadcast the vmap axis onto the operands in arg_idx that do not carry it,
  # so the flattening logic below sees uniformly batched operands.
  sizes = {args.shape[dim] for args, dim in zip(batched_args, batch_dims)
           if dim is not None}
  assert len(sizes) == 1, f"expected one vmap axis size, got {sizes}"
  axis_size, = sizes
  args, dims = list(batched_args), list(batch_dims)
  for i in arg_idx:
    if dims[i] is None:
      args[i] = jnp.broadcast_to(args[i][None], (axis_size,) + args[i].shape)
      dims[i] = 0
  return tuple(args), tuple(dims)

def _batcher_arg_idx(mask_type, num_args):
  # Operands that participate in batching; bias (index 3) is decided by the

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose operands so the vmapped axis is axis 0 before vmap, then transpose back: use in_axes=0 on stacked arrays
2. Replace vmap with einsum-based manual batching or jax.lax.map over axis 0
3. Re-express per-head operations without vmap (e.g. use the num_heads dimension natively instead of vmapping over it)

Example fix

# before
attn = jax.vmap(single_head_attention, in_axes=(1, 1, 1, None))(q, k, v, bias)  # bdims=1 -> error

# after
attn_h = jax.vmap(single_head_attention, in_axes=(0, 0, 0, None))(
    q.transpose(1, 0, 2), k.transpose(1, 0, 2), v.transpose(1, 0, 2), bias)
attn = attn_h.transpose(1, 0, 2)

def batch_dims_ok(bdims):
    return all(d in (0, None) for d in bdims)

def vmappable_attention_axes(shapes, in_axes) -> bool:
    return all(a in (0, None) for a in in_axes)

When it happens

Trigger: Using jax.vmap over dot_product_attention with in_axes pointing at a non-zero dimension of q/k/v/bias, or mixing axes such that an operand's batch dim is not 0 or None (e.g. vmap(fn, in_axes=(1, 1, 1))).

Common situations: Vectorizing attention over heads (in_axes=1 on BHT layout), over sequence, or transposing tensors then vmapping; multi-host code where vmap collapses onto unexpected axes.

Related errors

• {name} was requested to map a value of non-array type {core.
• {name} was requested to map its argument along axis {axis},
• {name} must have at least one non-None value in in_axes or a
• for vmap support, subclass {type(self)} must implement `batc
• under vmap, the {member_name} of {name} produced an output b

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