{"record":{"id":"efe92bd43372abde","repo":"jax-ml/jax","slug":"tpu-custom-call-does-not-support-non-trivial-batch","errorCode":null,"errorMessage":"tpu_custom_call does not support non-trivial batching.","messagePattern":"tpu_custom_call does not support non-trivial batching\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu_custom_call.py","lineNumber":95,"sourceCode":"  if (\n      ctx.is_forward_compat()\n      or backend is None\n      or not is_libtpu_at_least(\"0.0.47\")\n  ):\n    return _FWD_COMPAT_VERSION\n  if ir_version_override is not None:\n    return ir_version_override()\n  return None\n\n\ntpu_custom_call_p = core.Primitive(\"tpu_custom_call\")\ntpu_custom_call_p.multiple_results = True\ndispatch.simple_impl(tpu_custom_call_p)\n\n\ndef tpu_custom_call_batcher(axis_data, args, dims, **kwargs):\n  if axis_data.size != 1:\n    raise NotImplementedError(\n        \"tpu_custom_call does not support non-trivial batching.\"\n    )\n  unbatched_args = tuple(\n      a if (d is None or d is None) else a[d]\n      for a, d in zip(args, dims, strict=True)\n  )\n  out_unbatched = tpu_custom_call_p.bind(*unbatched_args, **kwargs)\n  out = tuple(o[None] for o in out_unbatched)\n  return out, (0,) * len(out)\nbatching.fancy_primitive_batchers[tpu_custom_call_p] = tpu_custom_call_batcher\n\n\nclass MemorySpace(enum.Enum):\n  HBM = enum.auto()\n  VMEM = enum.auto()\n  SEMAPHORE_MEM = enum.auto()\n  SMEM = enum.auto()\n  HOST = enum.auto()","sourceCodeStart":77,"sourceCodeEnd":113,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu_custom_call.py#L77-L113","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"# before\nout = jax.vmap(tpu_custom_call_fn)(xs)  # xs.shape[0] > 1\n# after\nout = tpu_custom_call_on_batch(xs)  # kernel handles leading batch dims itself","handlingStrategy":"fallback","validationCode":"if batch := xs.shape[0] > 1:\n    out = batched_tpu_custom_call(xs)   # kernel handles leading dims\nelse:\n    out = tpu_custom_call_fn(xs)","typeGuard":null,"tryCatchPattern":"try:\n    out = jax.vmap(fn)(xs)\nexcept NotImplementedError as e:\n    if 'non-trivial batching' in str(e): out = jax.lax.map(fn, xs)\n    else: raise","preventionTips":["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."],"tags":["jax","tpu","vmap","custom-call","not-implemented"],"backgroundTag":"vmap-unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}