jax-ml/jax · error · NotImplementedError

vmapping pallas_call with no arguments.

Error message

vmapping pallas_call with no arguments.

What it means

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.

Source

Thrown at jax/_src/pallas/pallas_call.py:434

    compiler_params: Any,
    cost_estimate: CostEstimate | None,
    out_avals: tuple[jax_core.AbstractValue, ...],
    metadata: FrozenDict[str, str] | None,
    name: str | None,
):
  """Batch the pallas_call by calling it in loop over the batch size.

  This function provides a fallback implementation of batching a pallas_call
  for the cases in which adding a batch dimension to the pallas grid is not
  supported. This is currently the case when the batched dimension corresponds
  to a dynamic axis or a scalar prefetch argument.

  This implementation builds a HLO loop that dynamic_slices the inputs according
  to the current iteration index and dynamic_updates an (initially empty) output
  allocation.
  """
  if not dims:
    raise NotImplementedError("vmapping pallas_call with no arguments.")

  (axis_size,) = {
      arg.shape[dim]
      for arg, dim in zip(args, dims)
      if dim is not None
  }

  args, dims = _broadcast_input_output_aliases(
      args,
      dims,
      input_output_aliases=input_output_aliases,
      axis_size=axis_size,
  )

  # The output arrays are completely overwritten, so we can just initialize
  # empty arrays.
  initial_state = [
      jnp.empty(tuple_insert(bm.array_aval.shape, 0, axis_size),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the batched arrays explicitly as arguments to pallas_call instead of closing over them in the kernel function
2. Check that at least one argument to the vmapped function has a non-None in_axis so dims is non-empty
3. If no batching is actually needed, remove the vmap wrapper
4. 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

Example fix

// before
k = pl.load_const(W)  # W captured, no args
closed = pallas_call(lambda: kernel())( )  # vmap(f) -> error
// after
f = jax.vmap(lambda x: pallas_call(kernel, out_shape=...)(x))
# x is a batched argument, so dims is non-empty

def has_batched_pallas_args(*args, in_axes) -> bool:
    return any(a is not None and ax is not None for a, ax in zip(args, in_axes))

try:
    jax.vmap(f)(xs)
except NotImplementedError as e:
    if 'no arguments' in str(e):
        # fold batch into grid or pass arrays explicitly

When it happens

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

Common situations: 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.

Related errors

• Batching over custom allocations is not supported yet.
• Batching over dynamic grid values is not supported yet.
• ragged_dot vmap over any dim but 0 - NYI
• reduce_window batching is not implemented for initial values
• pallas_call with a mesh does not support batching

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