jax-ml/jax · error · ValueError

Invalid axis {axis} for num_programs

Error message

Invalid axis {axis} for num_programs

What it means

get_num_programs was called with an axis index that doesn't map to any non-vmapped user grid axis. After skipping vmapped axes, the requested axis+1-th user axis was never found.

Source

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

    )
    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}"
      )
    # Lazily initialize accumulator offsets. This avoids the dependency on
    # `get_tpu_info()` when creating the `LoweringContext`.
    if self.accumulator_offsets is None:
      self.accumulator_offsets = [0] * info.num_mxus
    assert isinstance(aval.memory_space, tpu_core.AccMemorySpace)
    mxu_id = aval.memory_space.mxu_id
    assert 0 <= mxu_id < len(self.accumulator_offsets)
    base_entry = self.accumulator_offsets[mxu_id]
    num_rows = math.prod(aval.shape[:-1])
    assert num_rows % info.num_sublanes == 0

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure axis < rank of the user grid passed to the call
2. Match the axis argument to the actual grid rank (use num_programs(0) for 1-D grids)
3. Account for vmapped axes: the axis indexes user-visible axes only

Example fix

// before
num = pl.num_programs(1)  # grid is 1-D
// after
num = pl.num_programs(0)

user_grid_rank = len(grid) - num_vmapped_axes
assert 0 <= axis < user_grid_rank, f'num_programs axis {axis} out of range for grid {grid}'

When it happens

Trigger: Kernel calling pl.num_programs(axis) with axis >= number of user (non-vmapped) grid axes, e.g. num_programs(1) with a 1-D grid, or when vmap consumed the higher axes.

Common situations: Copied Triton-style code using num_programs(1) into a 1-D-grid Mosaic kernel; changing grid rank without updating axis args; vmapping over trailing grid axes.

Related errors

• program id was requested but no grid was provided.
• 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/60564f95b9f74c6b. Report an issue: GitHub.