{"record":{"id":"6e56bb12117460da","repo":"jax-ml/jax","slug":"failed-to-infer-the-output-layout-of-the-iota-ple","errorCode":null,"errorMessage":"Failed to infer the output layout of the iota. Please apply plgpu.layout_cast to its output right after its creation.","messagePattern":"Failed to infer the output layout of the iota\\. Please apply plgpu\\.layout_cast to its output right after its creation\\.","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":4687,"sourceCode":"    ctx: LoweringRuleContext, x, *, new_layout\n):\n  layout = new_layout.to_mgpu()\n  if ctx.avals_in[0].ndim == 0:  # scalar case\n    if layout != mgpu.WGSplatFragLayout():\n      raise ValueError(\n          \"Only plgpu.Layout.WG_SPLAT is supported for scalar values.\"\n      )\n    return x\n  return mgpu.dialect.layout_cast(x, mgpu.to_layout_attr(layout))\n\n\n@register_lowering_rule(lax.iota_p, mgpu.LoweringSemantics.Lane)\ndef _iota_lowering(\n    ctx: LoweringRuleContext, dtype, shape, dimension, sharding\n):\n  del sharding  # Unused.\n  if ctx.out_layout_hint is None:\n    raise RuntimeError(\n        \"Failed to infer the output layout of the iota. Please apply\"\n        \" plgpu.layout_cast to its output right after its creation.\"\n    )\n  mlir_dtype = mgpu_utils.dtype_to_ir_type(dtype)\n  is_signed = mgpu_utils.is_signed(dtype)\n  return mgpu.FragmentedArray.broadcasted_iota(\n      mlir_dtype, shape, dimension, ctx.out_layout_hint, is_signed=is_signed\n  )\n\n\n@register_lowering_rule(lax.iota_p, mgpu.LoweringSemantics.Warpgroup)\ndef _iota_lowering_wg(\n    ctx: LoweringRuleContext, dtype, shape, dimension, sharding\n):\n  del ctx, sharding\n  result_type = ir.VectorType.get(shape, mgpu_utils.dtype_to_ir_type(dtype))\n  return mgpu.dialect.broadcasted_iota(result_type, dimension)\n","sourceCodeStart":4669,"sourceCodeEnd":4705,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L4669-L4705","documentation":"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.","triggerScenarios":"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.","commonSituations":"Building index vectors for shuffles or gather addresses inside custom lane-level code in plgpu kernels.","solutions":["Wrap the iota output in plgpu.layout_cast immediately: plgpu.layout_cast(lax.iota(...), layout)","Compute the iota before entering the Lane-semantics region where a layout hint exists"],"exampleFix":"# before\nidx = lax.iota(np.int32, 32)  # in lane-level code\n# after\nidx = plgpu.layout_cast(lax.iota(np.int32, 32), plgpu.Layout.WARP_ROW_MAJOR)","handlingStrategy":"validation","validationCode":"i = lax.iota(np.int32, n)\ni = plgpu.layout_cast(i, my_layout)  # always cast iota output in lane-level code","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always layout_cast iota outputs immediately after creation in lane semantics"],"tags":["pallas","mosaic-gpu","iota","layout"],"backgroundTag":"missing-layout-hint","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}