{"record":{"id":"01552cebec5bf19b","repo":"jax-ml/jax","slug":"vmapping-pallas-call-with-no-arguments","errorCode":null,"errorMessage":"vmapping pallas_call with no arguments.","messagePattern":"vmapping pallas_call with no arguments\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/pallas_call.py","lineNumber":434,"sourceCode":"    compiler_params: Any,\n    cost_estimate: CostEstimate | None,\n    out_avals: tuple[jax_core.AbstractValue, ...],\n    metadata: FrozenDict[str, str] | None,\n    name: str | None,\n):\n  \"\"\"Batch the pallas_call by calling it in loop over the batch size.\n\n  This function provides a fallback implementation of batching a pallas_call\n  for the cases in which adding a batch dimension to the pallas grid is not\n  supported. This is currently the case when the batched dimension corresponds\n  to a dynamic axis or a scalar prefetch argument.\n\n  This implementation builds a HLO loop that dynamic_slices the inputs according\n  to the current iteration index and dynamic_updates an (initially empty) output\n  allocation.\n  \"\"\"\n  if not dims:\n    raise NotImplementedError(\"vmapping pallas_call with no arguments.\")\n\n  (axis_size,) = {\n      arg.shape[dim]\n      for arg, dim in zip(args, dims)\n      if dim is not None\n  }\n\n  args, dims = _broadcast_input_output_aliases(\n      args,\n      dims,\n      input_output_aliases=input_output_aliases,\n      axis_size=axis_size,\n  )\n\n  # The output arrays are completely overwritten, so we can just initialize\n  # empty arrays.\n  initial_state = [\n      jnp.empty(tuple_insert(bm.array_aval.shape, 0, axis_size),","sourceCodeStart":416,"sourceCodeEnd":452,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/pallas_call.py#L416-L452","documentation":"jax.pallas_call is a low-level Pallas kernel primitive whose batching rule (used by vmap) requires at least one batched argument with a concrete mapped dimension. When _batch_with_explicit_loop is entered with an empty dims sequence, there is no axis to map over, so JAX raises NotImplementedError instead of producing an ill-defined vmapped kernel. This mirrors the identical check in _pallas_call_batching_rule.","triggerScenarios":"Calling jax.vmap over a function whose traced body contains a pallas_call where every argument is unmapped (e.g. all inputs are constants closed over by the kernel, or the call takes no arguments, or vmap's in_axes are None for every operand).","commonSituations":"Wrapping a Pallas TPU/GPU kernel in vmap while passing scalars or module-level constants directly to pallas_call instead of as inputs; using in_axes=None everywhere; accidentally vmapping a function that only uses captured constants.","solutions":["Pass the batched arrays explicitly as arguments to pallas_call instead of closing over them in the kernel function","Check that at least one argument to the vmapped function has a non-None in_axis so dims is non-empty","If no batching is actually needed, remove the vmap wrapper","Rewrite the batched computation as an explicit loop (lax.fori_loop) or fold the batch dimension into the kernel grid instead of relying on vmap"],"exampleFix":"// before\nk = pl.load_const(W)  # W captured, no args\nclosed = pallas_call(lambda: kernel())( )  # vmap(f) -> error\n// after\nf = jax.vmap(lambda x: pallas_call(kernel, out_shape=...)(x))\n# x is a batched argument, so dims is non-empty","handlingStrategy":"validation","validationCode":"def has_batched_pallas_args(*args, in_axes) -> bool:\n    return any(a is not None and ax is not None for a, ax in zip(args, in_axes))","typeGuard":null,"tryCatchPattern":"try:\n    jax.vmap(f)(xs)\nexcept NotImplementedError as e:\n    if 'no arguments' in str(e):\n        # fold batch into grid or pass arrays explicitly","preventionTips":["Never close over arrays in Pallas kernels under vmap; always pass them as arguments","Keep at least one in_axis non-None for arrays consumed by pallas_call"],"tags":["jax","pallas","vmap","batching","not-implemented"],"backgroundTag":"vmap-unsupported-primitive","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}