jax-ml/jax · error · ValueError

program id was requested but no grid was provided.

Error message

program id was requested but no grid was provided.

What it means

Kernel lowering requested program_id (axis) but no user grid was provided. Program IDs index into the user-visible grid; with grid=None (e.g. fully vmapped or gridless kernel) there is no axis to return.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:249

  @contextlib.contextmanager
  def grid_name_context(self):
    # TODO(b/355036977): generalize this across other platforms
    if not self.grid_names:
      yield
      return
    grid_names = tuple(
        n for i, n in enumerate(self.grid_names) if i not in self.vmapped_dims
    )
    valid_grid_sizes = tuple(
        d for i, d in enumerate(self.grid_sizes) if i not in self.vmapped_dims
    )
    grid_env = zip(grid_names, valid_grid_sizes)
    with jax_core.extend_axis_env_nd(grid_env):
      yield

  def get_program_id(self, axis: int) -> ir.Value:
    if self.user_grid_indices is None:
      raise ValueError("program id was requested but no grid was provided.")
    return self.user_grid_indices[axis]

  def get_num_programs(self, axis: int) -> ir.Value:
    vmapped_axes = set(self.vmapped_dims)
    seen_user_axes = 0
    for i in range(self.grid_rank):
      seen_user_axes += int(i not in vmapped_axes)
      if seen_user_axes == axis + 1:
        break
    else:
      raise ValueError(f"Invalid axis {axis} for num_programs")
    return tpu.iteration_bound(i)

  def alloc_accumulator(self, aval: state.AbstractRef) -> AccRef:
    info = tpu_info.get_tpu_info()
    if not info.num_accumulators:
      raise ValueError(
          f"Accumulators are not available on TPU {info.chip_version}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a grid (e.g. grid=(1,) or the real iteration space) to the call that lowers the kernel
2. Remove the program_id usage if the kernel is meant to be gridless
3. If vmapping consumed all axes, restructure so a user grid remains

Example fix

// before
out = kernel(x)  # kernel uses pl.program_id(0), no grid
// after
out = pl.pallas_call(kernel, out_shape=..., grid=(1,))(x)

assert grid is not None, 'kernel body uses pl.program_id; a grid must be provided'

When it happens

Trigger: A kernel body calling pl.program_id(axis) while being lowered with no grid specified (grid=None) in the Mosaic call.

Common situations: Forgetting to pass grid= to the pallas/mosaic call; refactoring a gridless helper kernel to use program_id; vmapping over all grid axes leaving no user grid.

Understand the failure class

Background: Missing required parameter errors: what 'X is required' and 'the required X param is missing' mean, and how to fix them — this error's family across 27 libraries.

Related errors

• Invalid axis {axis} for num_programs
• Too many dynamic shapes in the input. Mosaic currently only
• Axis {axis} is out of bounds for grid {self.grid}
• Acc ref must be at least 2D, got shape {shape}
• Unsupported core type: {core_type}

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