jax-ml/jax · error · LookupError

Axis name {axis_name} not found in grid.

Error message

Axis name {axis_name} not found in grid.

What it means

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.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2628

    raise NotImplementedError
  return lower_jaxpr_to_triton_ir(ctx.context, jaxpr, ctx.block_infos, *args)


@register_lowering(ad_checkpoint.remat_p)
def _remat_lowering_rule(ctx: LoweringRuleContext, *args, jaxpr, **_):
  return lower_jaxpr_to_triton_ir(ctx.context, jaxpr, ctx.block_infos, *args)


triton_lowering_rules[ad_util.stop_gradient_p] = lambda _, x: x


@register_lowering(lax.axis_index_p)
def _axis_index_rule(ctx: LoweringRuleContext, *, axis_name: Hashable):
  grid_names = ctx.context.grid_mapping.grid_names
  if grid_names is not None and axis_name in grid_names:
    # We are querying a named axis corresponding to a grid dimension.
    return _program_id_lowering_rule(ctx, axis=grid_names.index(axis_name))
  raise LookupError(f"Axis name {axis_name} not found in grid.")


def _lower_jaxpr_to_for_loop(
    ctx: LoweringRuleContext,
    jaxpr: jax_core.Jaxpr,
    lower_bound,
    upper_bound,
    consts,
    *args,
    has_loop_index: bool,
    step: int = 1,
    bound_type: ir.IntegerType | None = None,
):
  if step != 1:
    raise NotImplementedError
  if bound_type is None or bound_type.width == 32:
    step_val = _i32_constant(step)
  else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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)
2. Use the program id directly instead: pl.program_id(axis) with the integer grid axis
3. Move axis_index calls outside the kernel and pass the index in as an input

Example fix

# before
pid = lax.axis_index('row')  # grid has no name 'row'

# after
pid = pl.program_id(0)

names = kernel.grid_mapping.grid_names if hasattr(kernel, 'grid_mapping') else None
assert names is not None and axis_name in names

def axis_in_grid(axis_name, grid_names) -> bool:
    return grid_names is not None and axis_name in grid_names

try:
    idx = lax.axis_index(name)
except LookupError:
    idx = pl.program_id(0)  # fallback to positional grid id

When it happens

Trigger: 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.

Common situations: Reusing axis names from pmap/sharded code inside a Mosaic kernel; forgetting to provide the named grid mapping when defining the kernel.

Related errors

• Axis {axis} is out of bounds for grid {self.grid}
• 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

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