jax-ml/jax · error · RuntimeError

Failed to infer the output layout of the iota. Please apply

Error message

Failed to infer the output layout of the iota. Please apply plgpu.layout_cast to its output right after its creation.

What it means

When lowering lax.iota at Lane semantics, the output layout cannot be inferred from context, so the developer must explicitly cast the iota's result to a layout immediately after creation. The backend surfaces this as a RuntimeError with that instruction.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:4687

    ctx: LoweringRuleContext, x, *, new_layout
):
  layout = new_layout.to_mgpu()
  if ctx.avals_in[0].ndim == 0:  # scalar case
    if layout != mgpu.WGSplatFragLayout():
      raise ValueError(
          "Only plgpu.Layout.WG_SPLAT is supported for scalar values."
      )
    return x
  return mgpu.dialect.layout_cast(x, mgpu.to_layout_attr(layout))


@register_lowering_rule(lax.iota_p, mgpu.LoweringSemantics.Lane)
def _iota_lowering(
    ctx: LoweringRuleContext, dtype, shape, dimension, sharding
):
  del sharding  # Unused.
  if ctx.out_layout_hint is None:
    raise RuntimeError(
        "Failed to infer the output layout of the iota. Please apply"
        " plgpu.layout_cast to its output right after its creation."
    )
  mlir_dtype = mgpu_utils.dtype_to_ir_type(dtype)
  is_signed = mgpu_utils.is_signed(dtype)
  return mgpu.FragmentedArray.broadcasted_iota(
      mlir_dtype, shape, dimension, ctx.out_layout_hint, is_signed=is_signed
  )


@register_lowering_rule(lax.iota_p, mgpu.LoweringSemantics.Warpgroup)
def _iota_lowering_wg(
    ctx: LoweringRuleContext, dtype, shape, dimension, sharding
):
  del ctx, sharding
  result_type = ir.VectorType.get(shape, mgpu_utils.dtype_to_ir_type(dtype))
  return mgpu.dialect.broadcasted_iota(result_type, dimension)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the iota output in plgpu.layout_cast immediately: plgpu.layout_cast(lax.iota(...), layout)
2. Compute the iota before entering the Lane-semantics region where a layout hint exists

Example fix

# before
idx = lax.iota(np.int32, 32)  # in lane-level code
# after
idx = plgpu.layout_cast(lax.iota(np.int32, 32), plgpu.Layout.WARP_ROW_MAJOR)

i = lax.iota(np.int32, n)
i = plgpu.layout_cast(i, my_layout)  # always cast iota output in lane-level code

When it happens

Trigger: Creating an iota/range inside a Lane-semantics context (e.g. warp-specialized or manually lowered code) where no layout hint propagates, then using it in layout-sensitive ops.

Common situations: Building index vectors for shuffles or gather addresses inside custom lane-level code in plgpu kernels.

Related errors

• packed cannot be specified if layout is specified.
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Swizzle {self.swizzle} requires the trailing dimension to be
• Can't instantiate {self} with arguments.
• Only TiledLayout supports reductions.

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