jax-ml/jax · error · NotImplementedError

pallas_call with a mesh does not support batching

Error message

pallas_call with a mesh does not support batching

What it means

Pallas kernels launched with a jax.sharding.Mesh (sharded/multi-device pallas_call via the mesh parameter) have no implemented batching rule. When vmap hits such a call, _pallas_call_batching_rule immediately raises NotImplementedError because combining vmap's batch dimension semantics with mesh-partitioned kernels is not supported.

Source

Thrown at jax/_src/pallas/pallas_call.py:542

  if all(bdim is None for bdim in dims):
    out = pallas_call_p.bind(
        *args,
        jaxpr=jaxpr,
        grid_mapping=grid_mapping,
        mesh=mesh,
        input_output_aliases=input_output_aliases,
        debug=debug,
        interpret=interpret,
        compiler_params=compiler_params,
        cost_estimate=cost_estimate,
        out_avals=out_avals,
        metadata=metadata,
        name=name,
    )
    return out, (None,) * len(out)

  if mesh is not None:
    raise NotImplementedError(
        "pallas_call with a mesh does not support batching"
    )

  def _maybe_squeeze_out_bdim(x: jax_typing.Array, bdim: int | batching.NotMapped
                              ) -> jax_typing.Array:
    return x if bdim is None else jnp.squeeze(x, axis=bdim)

  # this is the _global_ axis size if axis_data.explicit_mesh_axis is not None
  # we want to convert it to the local axis size
  axis_size = axis_data.size
  ema = axis_data.explicit_mesh_axis
  abs_mesh = get_abstract_mesh()
  if ema:
    mesh_size = math.prod(abs_mesh.shape[i] for i in ema)
    axis_size, ragged = divmod(axis_size, mesh_size)
    assert not ragged

  if axis_size == 1:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the vmap and fold the batch dimension into the pallas_call grid instead (e.g. add a grid axis of size batch)
2. Use a manual Python/lax.fori_loop over the batch dimension
3. Run one pallas_call per batch element via jax.lax.map with a non-vmapped body where feasible
4. Batch inside the kernel by declaring the batch dimension in BlockSpec/grid mappings rather than via vmap

Example fix

// before
f = jax.vmap(single_call)  # single_call uses pallas_call(..., mesh=mesh)
// after
def batched(x):
  grid = (x.shape[0], *orig_grid)  # batch as a grid axis
  return pallas_call(kernel, out_shape=out, grid=grid, mesh=mesh, ...)(x)

def can_vmap_pallas(call_kwargs) -> bool:
    return call_kwargs.get('mesh') is None

try:
    jax.vmap(f)(xs)
except NotImplementedError as e:
    if 'mesh' in str(e):
        # fold batch into the kernel grid instead

When it happens

Trigger: Applying jax.vmap to a function that calls pallas_call(..., mesh=mesh) with a non-None Mesh argument, e.g. multi-device TPU Pallas kernels driven by jax.lax.with_sharding_constraint or explicit mesh plumbing.

Common situations: Scaling a single-device Pallas kernel to a TPU pod with Mesh and forgetting a leftover vmap in the training/eval step; libraries (e.g. Mosaic-based attention kernels) that require meshes and user code that wraps them in vmap.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• Explicit sharding is not currently supported in the pallas-t
• {axis} mixes JAX mesh and Pallas mesh grid axes
• Batching over custom allocations is not supported yet.

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