{"record":{"id":"95ef4cded1102d94","repo":"jax-ml/jax","slug":"pallas-call-with-a-mesh-does-not-support-batching","errorCode":null,"errorMessage":"pallas_call with a mesh does not support batching","messagePattern":"pallas_call with a mesh does not support batching","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":542,"sourceCode":"  if all(bdim is None for bdim in dims):\n    out = pallas_call_p.bind(\n        *args,\n        jaxpr=jaxpr,\n        grid_mapping=grid_mapping,\n        mesh=mesh,\n        input_output_aliases=input_output_aliases,\n        debug=debug,\n        interpret=interpret,\n        compiler_params=compiler_params,\n        cost_estimate=cost_estimate,\n        out_avals=out_avals,\n        metadata=metadata,\n        name=name,\n    )\n    return out, (None,) * len(out)\n\n  if mesh is not None:\n    raise NotImplementedError(\n        \"pallas_call with a mesh does not support batching\"\n    )\n\n  def _maybe_squeeze_out_bdim(x: jax_typing.Array, bdim: int | batching.NotMapped\n                              ) -> jax_typing.Array:\n    return x if bdim is None else jnp.squeeze(x, axis=bdim)\n\n  # this is the _global_ axis size if axis_data.explicit_mesh_axis is not None\n  # we want to convert it to the local axis size\n  axis_size = axis_data.size\n  ema = axis_data.explicit_mesh_axis\n  abs_mesh = get_abstract_mesh()\n  if ema:\n    mesh_size = math.prod(abs_mesh.shape[i] for i in ema)\n    axis_size, ragged = divmod(axis_size, mesh_size)\n    assert not ragged\n\n  if axis_size == 1:","sourceCodeStart":524,"sourceCodeEnd":560,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L524-L560","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Remove the vmap and fold the batch dimension into the pallas_call grid instead (e.g. add a grid axis of size batch)","Use a manual Python/lax.fori_loop over the batch dimension","Run one pallas_call per batch element via jax.lax.map with a non-vmapped body where feasible","Batch inside the kernel by declaring the batch dimension in BlockSpec/grid mappings rather than via vmap"],"exampleFix":"// before\nf = jax.vmap(single_call)  # single_call uses pallas_call(..., mesh=mesh)\n// after\ndef batched(x):\n  grid = (x.shape[0], *orig_grid)  # batch as a grid axis\n  return pallas_call(kernel, out_shape=out, grid=grid, mesh=mesh, ...)(x)","handlingStrategy":"validation","validationCode":"def can_vmap_pallas(call_kwargs) -> bool:\n    return call_kwargs.get('mesh') is None","typeGuard":null,"tryCatchPattern":"try:\n    jax.vmap(f)(xs)\nexcept NotImplementedError as e:\n    if 'mesh' in str(e):\n        # fold batch into the kernel grid instead","preventionTips":["Reserve vmap for mesh-free pallas calls","Express batching via grid/BlockSpec axes on multi-device kernels"],"tags":["jax","pallas","vmap","mesh","sharding","not-implemented"],"backgroundTag":"vmap-unsupported-primitive","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}