{"record":{"id":"12b2dcb0d99708b0","repo":"jax-ml/jax","slug":"program-id-was-requested-but-no-grid-was-provided","errorCode":null,"errorMessage":"program id was requested but no grid was provided.","messagePattern":"program id was requested but no grid was provided\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":249,"sourceCode":"  @contextlib.contextmanager\n  def grid_name_context(self):\n    # TODO(b/355036977): generalize this across other platforms\n    if not self.grid_names:\n      yield\n      return\n    grid_names = tuple(\n        n for i, n in enumerate(self.grid_names) if i not in self.vmapped_dims\n    )\n    valid_grid_sizes = tuple(\n        d for i, d in enumerate(self.grid_sizes) if i not in self.vmapped_dims\n    )\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}\"","sourceCodeStart":231,"sourceCodeEnd":267,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L231-L267","documentation":"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.","triggerScenarios":"A kernel body calling pl.program_id(axis) while being lowered with no grid specified (grid=None) in the Mosaic call.","commonSituations":"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.","solutions":["Pass a grid (e.g. grid=(1,) or the real iteration space) to the call that lowers the kernel","Remove the program_id usage if the kernel is meant to be gridless","If vmapping consumed all axes, restructure so a user grid remains"],"exampleFix":"// before\nout = kernel(x)  # kernel uses pl.program_id(0), no grid\n// after\nout = pl.pallas_call(kernel, out_shape=..., grid=(1,))(x)","handlingStrategy":"validation","validationCode":"assert grid is not None, 'kernel body uses pl.program_id; a grid must be provided'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass grid= to pallas/mosaic calls whose kernels use program_id","Pass grid=(1,) for gridless kernels that still query program_id"],"tags":["jax","pallas","mosaic","grid","program-id","lowering"],"backgroundTag":"missing-required-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}