jax-ml/jax · error · ValueError

Axis {axis} is out of bounds for grid {self.grid}

Error message

Axis {axis} is out of bounds for grid {self.grid}

What it means

Kernel grids are tuples; GridContext.size(axis) (used by program_id and num_programs primitives) indexes the grid and converts a caught IndexError into ValueError with a clear message. Requesting an axis >= grid rank (or a too-negative axis) triggers it.

Source

Thrown at jax/_src/pallas/core.py:341

  def name(self) -> Any:
    return f"{self.memory_space}@{self.mesh.core_type.name}"

  @property
  def memory_kind(self) -> str:
    return jax_core.mem_space_to_kind(self.memory_space)


@dataclasses.dataclass(frozen=True)
class PallasGridContext:
  grid: GridMappingGrid
  mapped_dims: tuple[int, ...]

  def size(self, axis: int) -> int | DynamicGridDim:
    valid_grid = tuple(self.grid)
    try:
      size = valid_grid[axis]
    except IndexError as e:
      raise ValueError(
          f"Axis {axis} is out of bounds for grid {self.grid}"
      ) from e
    return size


@dataclasses.dataclass
class PallasTracingEnv(threading.local):
  grid_context: PallasGridContext | None = None
  grid_env_stack: list[GridEnv] = dataclasses.field(default_factory=list)
  is_interpret_mode: bool = False
  dynamic_shapes: bool = False
  module_export_fn: Callable[[mlir.ir.Module], None] | None = None

_pallas_tracing_env = PallasTracingEnv()


def axis_frame() -> PallasGridContext:
  # This is like jax_core.axis_frame, except there should only ever be one

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make len(grid) match the highest axis passed to program_id/num_programs
2. Derive grid programmatically: grid = tuple(cdiv(a, b) for a, b in zip(shape, block_shape))
3. Double-check the grid argument of pallas_call, not just the kernel body

Example fix

# before
kernel = pl.pallas_call(f, out_shape, grid=(8, 4))  # f uses pl.program_id(2)
# after
kernel = pl.pallas_call(f, out_shape, grid=(8, 4, 1))  # or drop program_id(2)

assert -len(grid) <= axis < len(grid), f'axis {axis} out of bounds for grid {grid}'

When it happens

Trigger: Calling pl.program_id(axis) or pl.num_programs(axis) with axis >= len(grid), e.g. a 2D grid (8, 4) and axis=2; computing axis from array ndim while the grid has fewer dims.

Common situations: Generalizing a kernel from 1D to 2D blocks but forgetting to extend the grid; deriving axis programmatically (axis = x.ndim) that outranks the pallas_call grid; off-by-one in loop bounds.

Related errors

• Revisited block {output_ranges[i]} of output {i} in iteratio
• Dimensions with parallel semantics must form a prefix of the
• program id was requested but no grid was provided.
• Invalid axis {axis} for num_programs
• Non-trivial windowing is not supported for grid-free pallas_

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