{"record":{"id":"9d8e4f09adb596f8","repo":"jax-ml/jax","slug":"batching-over-dynamic-grid-values-is-not-supported","errorCode":null,"errorMessage":"Batching over dynamic grid values is not supported yet.","messagePattern":"Batching over dynamic grid values is not supported yet\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":2696,"sourceCode":"  dimension_semantics = (PARALLEL,) + dimension_semantics\n  all_args: EmitPipelinePrimitiveArgs = args_tree.unflatten(args_flat)\n\n  _, dynamic_dims, _, _, flat_ref_dims, alloc_dims = jax_util.split_list(dims, [\n      len(all_args.all_index_map_consts),\n      len(all_args.dynamic_grid_spec),\n      int(all_args.has_core_id),\n      len(all_args.body_consts),\n      len(all_args.refs_flat)])\n\n  if any(d is not None for d in alloc_dims):\n    raise NotImplementedError(\n        \"Batching over custom allocations is not supported yet.\"\n    )\n\n  batch_size = axis_data.size\n\n  if any(d is not None for d in dynamic_dims):\n    raise NotImplementedError(\n      \"Batching over dynamic grid values is not supported yet.\")\n\n  batched_block_mappings = map(\n      functools.partial(_batch_block_mapping, grid_mapping, batch_size),\n      map(_ref_to_value_aval, all_args.refs_flat),\n      flat_ref_dims, grid_mapping.block_mappings)\n\n  index_map_tree_args, index_map_tree_kwargs = (\n      grid_mapping.index_map_tree.unflatten(grid_mapping.index_map_avals))\n  assert not index_map_tree_kwargs\n  batched_index_map_args = (\n      pallas_core.index_map_grid_aval, *index_map_tree_args)\n  batched_index_map_avals, batched_index_map_tree = tree_util.tree_flatten(\n      (batched_index_map_args, {}))\n\n  axis_size_is_dynamic = not isinstance(batch_size, int)\n  new_grid_dim = (pallas_core.dynamic_grid_dim\n                  if axis_size_is_dynamic else batch_size)","sourceCodeStart":2678,"sourceCodeEnd":2714,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L2678-L2714","documentation":"Raised by the Mosaic Pallas pipeline when vmap-style batching is applied to a kernel whose grid depends on dynamic (traced) values. The batching transform can replicate block mappings for a known batch size, but it cannot reconstruct index maps when grid extents themselves are computed at trace time, so it refuses.","triggerScenarios":"Calling jax.vmap (or another batching transform) over a function that invokes a pallas_call/mosaic kernel whose grid tuple contains dynamic values (e.g. grid=(x.shape[0],) where the shape depends on the batched axis, or a grid built from a traced computation).","commonSituations":"Batching a TPU Pallas kernel that computes its grid from an input whose leading dim is being vmapped; stacking per-example kernels where grid=(num_examples,) and num_examples is derived from the batch.","solutions":["Make the grid static: compute grid extents from Python ints or from shapes that do not depend on the batched axis (e.g. hardcode the per-example grid and let vmap add a leading grid dimension)","Restructure so the batch appears as an extra grid dimension that the kernel indexes via program_id instead of a dynamic grid value","Move batching inside the kernel body (manual loop over the batch within one launch) instead of using vmap on the outside","Check for a newer JAX version — batching support for Pallas is actively expanding"],"exampleFix":"# before\ngrid = (x.shape[0] // 128,)  # x.shape[0] depends on vmapped axis\nkern = pallas_call(fn, out_shape=..., grid=grid)\nout = jax.vmap(lambda x: kern(x))(xs)\n\n# after\ndef fn(x_ref, o_ref):\n  ...\n# static per-example grid; batch becomes grid dim\nkern = pallas_call(fn, out_shape=..., grid=(xs.shape[0], xs.shape[1] // 128))\nout = kern(xs)","handlingStrategy":"validation","validationCode":"def check_grid_static(kernel_fn, *args, **kwargs):\n  import jax\n  try:\n    jax.make_jaxpr(lambda *a: kernel_fn(*a, **kwargs))(*args)\n  except Exception as e:\n    raise RuntimeError(f\"kernel trace failed (dynamic grid?)\") from e\n# ensure grid extents are Python ints:\nassert all(isinstance(g, int) for g in grid), f\"dynamic grid: {grid}\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep pallas_call grids as tuples of Python ints computed from static shapes","Never derive grid extents from the axis being vmapped","Add unit tests that vmap your kernel entry points in CI"],"tags":["jax","pallas","mosaic","vmap","batching","tpu","not-implemented"],"backgroundTag":"vmap-unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}