jax-ml/jax · error · NotImplementedError

Batching over custom allocations is not supported yet.

Error message

Batching over custom allocations is not supported yet.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/pipeline.py:2689

emit_pipeline_p.to_lojax = _emit_pipeline_to_lojax


def _emit_pipeline_batching_rule(
    axis_data, args_flat, dims, *, grid_mapping, dimension_semantics, args_tree,
    **params
):
  dimension_semantics = (PARALLEL,) + dimension_semantics
  all_args: EmitPipelinePrimitiveArgs = args_tree.unflatten(args_flat)

  _, dynamic_dims, _, _, flat_ref_dims, alloc_dims = jax_util.split_list(dims, [
      len(all_args.all_index_map_consts),
      len(all_args.dynamic_grid_spec),
      int(all_args.has_core_id),
      len(all_args.body_consts),
      len(all_args.refs_flat)])

  if any(d is not None for d in alloc_dims):
    raise NotImplementedError(
        "Batching over custom allocations is not supported yet."
    )

  batch_size = axis_data.size

  if any(d is not None for d in dynamic_dims):
    raise NotImplementedError(
      "Batching over dynamic grid values is not supported yet.")

  batched_block_mappings = map(
      functools.partial(_batch_block_mapping, grid_mapping, batch_size),
      map(_ref_to_value_aval, all_args.refs_flat),
      flat_ref_dims, grid_mapping.block_mappings)

  index_map_tree_args, index_map_tree_kwargs = (
      grid_mapping.index_map_tree.unflatten(grid_mapping.index_map_avals))
  assert not index_map_tree_kwargs
  batched_index_map_args = (

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
batched_kernel = jax.vmap(kernel, in_axes=0)  # kernel uses custom allocations
# after
def kernel_batched(x):
  b = x.shape[0]
  return emit_pipeline(body, grid=(b, ...), in_specs=..., out_specs=...)(x)

Defensive patterns

Strategy: fallback

Validate before calling

import jax
try:
    jax.vmap(kernel)(x)
except NotImplementedError:
    ...  # detect before call only by knowing allocations are custom; assert in helper
assert not has_custom_allocations(kernel), 'vmap unsupported over custom allocations'

Try / catch

try:
    out = jax.vmap(kernel)(x)
except NotImplementedError as e:
    if 'custom allocations' in str(e):
        out = jax.lax.map(kernel, x)  # sequential fallback

Prevention

Put the batch dimension into the grid instead of vmapping pipelined kernels
Avoid custom-sized allocations when you plan to vmap

When it happens

Trigger: 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.

Common situations: Wrapping a Pallas pipeline kernel in vmap for batched inference/training; upgrading kernels that allocated scratch by a per-example size.

Related errors

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