jax-ml/jax · error · NotImplementedError

tpu_custom_call does not support non-trivial batching.

Error message

tpu_custom_call does not support non-trivial batching.

What it means

tpu_custom_call does not support vmapping over a batch dimension of size > 1; its batching rule only handles a symbolic/size-1 batch axis. Any non-trivial vmap of a tpu_custom_call raises NotImplementedError.

Source

Thrown at jax/_src/tpu_custom_call.py:95

  if (
      ctx.is_forward_compat()
      or backend is None
      or not is_libtpu_at_least("0.0.47")
  ):
    return _FWD_COMPAT_VERSION
  if ir_version_override is not None:
    return ir_version_override()
  return None


tpu_custom_call_p = core.Primitive("tpu_custom_call")
tpu_custom_call_p.multiple_results = True
dispatch.simple_impl(tpu_custom_call_p)


def tpu_custom_call_batcher(axis_data, args, dims, **kwargs):
  if axis_data.size != 1:
    raise NotImplementedError(
        "tpu_custom_call does not support non-trivial batching."
    )
  unbatched_args = tuple(
      a if (d is None or d is None) else a[d]
      for a, d in zip(args, dims, strict=True)
  )
  out_unbatched = tpu_custom_call_p.bind(*unbatched_args, **kwargs)
  out = tuple(o[None] for o in out_unbatched)
  return out, (0,) * len(out)
batching.fancy_primitive_batchers[tpu_custom_call_p] = tpu_custom_call_batcher


class MemorySpace(enum.Enum):
  HBM = enum.auto()
  VMEM = enum.auto()
  SEMAPHORE_MEM = enum.auto()
  SMEM = enum.auto()
  HOST = enum.auto()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Instead of vmap, fold the batch into the leading dims of a single custom call (many TPU custom calls already batch over leading dims internally).
Write the kernel to handle a batch dimension explicitly and call it once on the stacked input.
If vmap is unavoidable, use vmap with axis_size 1 or loop/map manually over the batch with lax.map / python loop.

Example fix

# before
out = jax.vmap(tpu_custom_call_fn)(xs)  # xs.shape[0] > 1
# after
out = tpu_custom_call_on_batch(xs)  # kernel handles leading batch dims itself

Defensive patterns

Strategy: fallback

Validate before calling

if batch := xs.shape[0] > 1:
    out = batched_tpu_custom_call(xs)   # kernel handles leading dims
else:
    out = tpu_custom_call_fn(xs)

Try / catch

try:
    out = jax.vmap(fn)(xs)
except NotImplementedError as e:
    if 'non-trivial batching' in str(e): out = jax.lax.map(fn, xs)
    else: raise

Prevention

Prefer folding batch dims into the custom call itself over vmap.
Use jax.lax.map for sequential batching of TPU custom calls.
Avoid vmap around Pallas/TPU kernels unless axis_size == 1.

When it happens

Trigger: jax.vmap over a function containing a tpu_custom_call where the mapped axis size != 1, e.g. vmap(fn)(batched_inputs) with batch size >= 2.

Common situations: Using Pallas/Mosaic TPU kernels or custom call wrappers inside vmap or a batched jit; code that worked scalar becomes batched in a training loop.

Related errors

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