{"record":{"id":"f6f2229229c1845f","repo":"jax-ml/jax","slug":"currently-only-support-batch-dim-in-0-none-but","errorCode":null,"errorMessage":"Currently only support batch_dim in [0, None], but got {dim=}","messagePattern":"Currently only support batch_dim in \\[0, None\\], but got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/cudnn/fused_attention_stablehlo.py","lineNumber":805,"sourceCode":"    backend_config=backend_config,\n    operand_layouts=default_layouts(\n      *[ir.RankedTensorType(operand.type).shape for operand in operands]),\n    result_layouts=result_layouts,\n  )\n  dqkv = (hlo.transpose(out.results[0], grad_transpose_perm),\n          hlo.transpose(out.results[1], grad_transpose_perm),\n          hlo.transpose(out.results[2], grad_transpose_perm))\n  # Only keep dQ, dK, dV and dBias here\n  if has_dbias:\n    return dqkv + (out.results[3],)\n  else:\n    return dqkv\n\n# batcher\ndef _check_valid_batch_dims(bdims):\n  for dim in bdims:\n    if dim not in [0, None]:\n      raise NotImplementedError(\n        f\"Currently only support batch_dim in [0, None], but got {dim=}\")\n\ndef _broadcast_unbatched_args(batched_args, batch_dims, arg_idx):\n  # Broadcast the vmap axis onto the operands in arg_idx that do not carry it,\n  # so the flattening logic below sees uniformly batched operands.\n  sizes = {args.shape[dim] for args, dim in zip(batched_args, batch_dims)\n           if dim is not None}\n  assert len(sizes) == 1, f\"expected one vmap axis size, got {sizes}\"\n  axis_size, = sizes\n  args, dims = list(batched_args), list(batch_dims)\n  for i in arg_idx:\n    if dims[i] is None:\n      args[i] = jnp.broadcast_to(args[i][None], (axis_size,) + args[i].shape)\n      dims[i] = 0\n  return tuple(args), tuple(dims)\n\ndef _batcher_arg_idx(mask_type, num_args):\n  # Operands that participate in batching; bias (index 3) is decided by the","sourceCodeStart":787,"sourceCodeEnd":823,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/cudnn/fused_attention_stablehlo.py#L787-L823","documentation":"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.","triggerScenarios":"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))).","commonSituations":"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.","solutions":["Transpose operands so the vmapped axis is axis 0 before vmap, then transpose back: use in_axes=0 on stacked arrays","Replace vmap with einsum-based manual batching or jax.lax.map over axis 0","Re-express per-head operations without vmap (e.g. use the num_heads dimension natively instead of vmapping over it)"],"exampleFix":"# before\nattn = jax.vmap(single_head_attention, in_axes=(1, 1, 1, None))(q, k, v, bias)  # bdims=1 -> error\n\n# after\nattn_h = jax.vmap(single_head_attention, in_axes=(0, 0, 0, None))(\n    q.transpose(1, 0, 2), k.transpose(1, 0, 2), v.transpose(1, 0, 2), bias)\nattn = attn_h.transpose(1, 0, 2)","handlingStrategy":"validation","validationCode":"def batch_dims_ok(bdims):\n    return all(d in (0, None) for d in bdims)","typeGuard":"def vmappable_attention_axes(shapes, in_axes) -> bool:\n    return all(a in (0, None) for a in in_axes)","tryCatchPattern":null,"preventionTips":["Conventionalize vmapped operands to axis 0 (stack, don't transpose-then-vmap)","Prefer native num_heads batching or einsum over vmap for attention heads"],"tags":["jax","vmap","cudnn","batching","in-axes"],"backgroundTag":"vmap-unsupported-axis","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}