{"record":{"id":"29f42953706bbd94","repo":"jax-ml/jax","slug":"batching-over-custom-allocations-is-not-supported","errorCode":null,"errorMessage":"Batching over custom allocations is not supported yet.","messagePattern":"Batching over custom allocations is not supported yet\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/pipeline.py","lineNumber":2689,"sourceCode":"emit_pipeline_p.to_lojax = _emit_pipeline_to_lojax\n\n\ndef _emit_pipeline_batching_rule(\n    axis_data, args_flat, dims, *, grid_mapping, dimension_semantics, args_tree,\n    **params\n):\n  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 = (","sourceCodeStart":2671,"sourceCodeEnd":2707,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/pipeline.py#L2671-L2707","documentation":"The pipeline batching rule (vmap over emit_pipeline) supports batching over dynamic dims, refs and body constants, but not over allocations that have custom allocation dimensions (alloc_dims). If any allocation dim is batched, NotImplementedError is raised.","triggerScenarios":"Applying jax.vmap to a function that calls a pipelined kernel whose allocations block defines custom-sized scratch buffers, so the allocation itself would need a new leading batch dimension.","commonSituations":"Wrapping a Pallas pipeline kernel in vmap for batched inference/training; upgrading kernels that allocated scratch by a per-example size.","solutions":["Move the batch dimension into the grid instead: iterate over examples with an extra grid axis and corresponding BlockSpec dim","Manually vmap by looping/stacking over the batch outside the pipeline call","Use dynamic grid dimensions (jax.Array grid entries) instead of custom allocations for batch-varying sizes"],"exampleFix":"# before\nbatched_kernel = jax.vmap(kernel, in_axes=0)  # kernel uses custom allocations\n# after\ndef kernel_batched(x):\n  b = x.shape[0]\n  return emit_pipeline(body, grid=(b, ...), in_specs=..., out_specs=...)(x)","handlingStrategy":"fallback","validationCode":"import jax\ntry:\n    jax.vmap(kernel)(x)\nexcept NotImplementedError:\n    ...  # detect before call only by knowing allocations are custom; assert in helper\nassert not has_custom_allocations(kernel), 'vmap unsupported over custom allocations'","typeGuard":null,"tryCatchPattern":"try:\n    out = jax.vmap(kernel)(x)\nexcept NotImplementedError as e:\n    if 'custom allocations' in str(e):\n        out = jax.lax.map(kernel, x)  # sequential fallback","preventionTips":["Put the batch dimension into the grid instead of vmapping pipelined kernels","Avoid custom-sized allocations when you plan to vmap"],"tags":["jax","pallas","vmap","batching","allocations","not-implemented"],"backgroundTag":"vmap-unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}