{"record":{"id":"60564f95b9f74c6b","repo":"jax-ml/jax","slug":"invalid-axis-axis-for-num-programs","errorCode":null,"errorMessage":"Invalid axis {axis} for num_programs","messagePattern":"Invalid axis (.+?) for num_programs","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":260,"sourceCode":"    )\n    grid_env = zip(grid_names, valid_grid_sizes)\n    with jax_core.extend_axis_env_nd(grid_env):\n      yield\n\n  def get_program_id(self, axis: int) -> ir.Value:\n    if self.user_grid_indices is None:\n      raise ValueError(\"program id was requested but no grid was provided.\")\n    return self.user_grid_indices[axis]\n\n  def get_num_programs(self, axis: int) -> ir.Value:\n    vmapped_axes = set(self.vmapped_dims)\n    seen_user_axes = 0\n    for i in range(self.grid_rank):\n      seen_user_axes += int(i not in vmapped_axes)\n      if seen_user_axes == axis + 1:\n        break\n    else:\n      raise ValueError(f\"Invalid axis {axis} for num_programs\")\n    return tpu.iteration_bound(i)\n\n  def alloc_accumulator(self, aval: state.AbstractRef) -> AccRef:\n    info = tpu_info.get_tpu_info()\n    if not info.num_accumulators:\n      raise ValueError(\n          f\"Accumulators are not available on TPU {info.chip_version}\"\n      )\n    # Lazily initialize accumulator offsets. This avoids the dependency on\n    # `get_tpu_info()` when creating the `LoweringContext`.\n    if self.accumulator_offsets is None:\n      self.accumulator_offsets = [0] * info.num_mxus\n    assert isinstance(aval.memory_space, tpu_core.AccMemorySpace)\n    mxu_id = aval.memory_space.mxu_id\n    assert 0 <= mxu_id < len(self.accumulator_offsets)\n    base_entry = self.accumulator_offsets[mxu_id]\n    num_rows = math.prod(aval.shape[:-1])\n    assert num_rows % info.num_sublanes == 0","sourceCodeStart":242,"sourceCodeEnd":278,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L242-L278","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Ensure axis < rank of the user grid passed to the call","Match the axis argument to the actual grid rank (use num_programs(0) for 1-D grids)","Account for vmapped axes: the axis indexes user-visible axes only"],"exampleFix":"// before\nnum = pl.num_programs(1)  # grid is 1-D\n// after\nnum = pl.num_programs(0)","handlingStrategy":"validation","validationCode":"user_grid_rank = len(grid) - num_vmapped_axes\nassert 0 <= axis < user_grid_rank, f'num_programs axis {axis} out of range for grid {grid}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Match the axis argument of num_programs to the actual (non-vmapped) grid rank","When changing grid rank, grep the kernel for hardcoded axis indices"],"tags":["jax","pallas","mosaic","grid","axis","lowering"],"backgroundTag":"axis-out-of-range","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}