{"record":{"id":"8baf0e044d72b15a","repo":"jax-ml/jax","slug":"axis-name-axis-name-not-found-in-grid","errorCode":null,"errorMessage":"Axis name {axis_name} not found in grid.","messagePattern":"Axis name (.+?) not found in grid\\.","errorType":"exception","errorClass":"LookupError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":2628,"sourceCode":"    raise NotImplementedError\n  return lower_jaxpr_to_triton_ir(ctx.context, jaxpr, ctx.block_infos, *args)\n\n\n@register_lowering(ad_checkpoint.remat_p)\ndef _remat_lowering_rule(ctx: LoweringRuleContext, *args, jaxpr, **_):\n  return lower_jaxpr_to_triton_ir(ctx.context, jaxpr, ctx.block_infos, *args)\n\n\ntriton_lowering_rules[ad_util.stop_gradient_p] = lambda _, x: x\n\n\n@register_lowering(lax.axis_index_p)\ndef _axis_index_rule(ctx: LoweringRuleContext, *, axis_name: Hashable):\n  grid_names = ctx.context.grid_mapping.grid_names\n  if grid_names is not None and axis_name in grid_names:\n    # We are querying a named axis corresponding to a grid dimension.\n    return _program_id_lowering_rule(ctx, axis=grid_names.index(axis_name))\n  raise LookupError(f\"Axis name {axis_name} not found in grid.\")\n\n\ndef _lower_jaxpr_to_for_loop(\n    ctx: LoweringRuleContext,\n    jaxpr: jax_core.Jaxpr,\n    lower_bound,\n    upper_bound,\n    consts,\n    *args,\n    has_loop_index: bool,\n    step: int = 1,\n    bound_type: ir.IntegerType | None = None,\n):\n  if step != 1:\n    raise NotImplementedError\n  if bound_type is None or bound_type.width == 32:\n    step_val = _i32_constant(step)\n  else:","sourceCodeStart":2610,"sourceCodeEnd":2646,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L2610-L2646","documentation":"lax.axis_index(axis_name) inside a Pallas Triton kernel only knows axis names that correspond to named dimensions of the kernel's grid (grid_mapping.grid_names). If the name isn't among the grid names (or grid_names is None), lowering raises LookupError.","triggerScenarios":"Calling lax.axis_index('rows') in a kernel launched with a plain Grid without name_map, or using a collectives axis name that exists only in the outer jit/pmap context, not in the pallas grid.","commonSituations":"Reusing axis names from pmap/sharded code inside a Mosaic kernel; forgetting to provide the named grid mapping when defining the kernel.","solutions":["Launch the kernel with a grid mapping that names its dimensions so axis_name matches (pass grid with names / name mapping used by your pallas version)","Use the program id directly instead: pl.program_id(axis) with the integer grid axis","Move axis_index calls outside the kernel and pass the index in as an input"],"exampleFix":"# before\npid = lax.axis_index('row')  # grid has no name 'row'\n\n# after\npid = pl.program_id(0)","handlingStrategy":"validation","validationCode":"names = kernel.grid_mapping.grid_names if hasattr(kernel, 'grid_mapping') else None\nassert names is not None and axis_name in names","typeGuard":"def axis_in_grid(axis_name, grid_names) -> bool:\n    return grid_names is not None and axis_name in grid_names","tryCatchPattern":"try:\n    idx = lax.axis_index(name)\nexcept LookupError:\n    idx = pl.program_id(0)  # fallback to positional grid id","preventionTips":["Prefer pl.program_id(i) over named axes in Mosaic kernels","Define grid names at kernel definition and keep them in sync with axis_index calls"],"tags":["jax","pallas","triton","grid","axis-index"],"backgroundTag":"axis-name-not-found","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}