{"record":{"id":"9f914eb7579b2f1e","repo":"jax-ml/jax","slug":"axis-axis-is-out-of-bounds-for-grid-self-grid","errorCode":null,"errorMessage":"Axis {axis} is out of bounds for grid {self.grid}","messagePattern":"Axis (.+?) is out of bounds for grid (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/core.py","lineNumber":341,"sourceCode":"  def name(self) -> Any:\n    return f\"{self.memory_space}@{self.mesh.core_type.name}\"\n\n  @property\n  def memory_kind(self) -> str:\n    return jax_core.mem_space_to_kind(self.memory_space)\n\n\n@dataclasses.dataclass(frozen=True)\nclass PallasGridContext:\n  grid: GridMappingGrid\n  mapped_dims: tuple[int, ...]\n\n  def size(self, axis: int) -> int | DynamicGridDim:\n    valid_grid = tuple(self.grid)\n    try:\n      size = valid_grid[axis]\n    except IndexError as e:\n      raise ValueError(\n          f\"Axis {axis} is out of bounds for grid {self.grid}\"\n      ) from e\n    return size\n\n\n@dataclasses.dataclass\nclass PallasTracingEnv(threading.local):\n  grid_context: PallasGridContext | None = None\n  grid_env_stack: list[GridEnv] = dataclasses.field(default_factory=list)\n  is_interpret_mode: bool = False\n  dynamic_shapes: bool = False\n  module_export_fn: Callable[[mlir.ir.Module], None] | None = None\n\n_pallas_tracing_env = PallasTracingEnv()\n\n\ndef axis_frame() -> PallasGridContext:\n  # This is like jax_core.axis_frame, except there should only ever be one","sourceCodeStart":323,"sourceCodeEnd":359,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/core.py#L323-L359","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Make len(grid) match the highest axis passed to program_id/num_programs","Derive grid programmatically: grid = tuple(cdiv(a, b) for a, b in zip(shape, block_shape))","Double-check the grid argument of pallas_call, not just the kernel body"],"exampleFix":"# before\nkernel = pl.pallas_call(f, out_shape, grid=(8, 4))  # f uses pl.program_id(2)\n# after\nkernel = pl.pallas_call(f, out_shape, grid=(8, 4, 1))  # or drop program_id(2)","handlingStrategy":"validation","validationCode":"assert -len(grid) <= axis < len(grid), f'axis {axis} out of bounds for grid {grid}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive grid from block_shape: grid = tuple(-(-a // b) for a, b in zip(shape, blocks))","Keep axis usage consistent with grid rank in kernel templates"],"tags":["jax","pallas","grid","index-out-of-bounds"],"backgroundTag":"index-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}