{"record":{"id":"876dc94f57050bae","repo":"jax-ml/jax","slug":"relayout-p-is-not-supported-with-lane-semantics","errorCode":null,"errorMessage":"relayout_p is not supported with Lane semantics.","messagePattern":"relayout_p is not supported with Lane semantics\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":4635,"sourceCode":"  exception = jax.tree.unflatten(exception_tree, payload)\n  assert isinstance(exception, checkify.FailedCheckError)\n\n  # check_p has an inverted predicate compared to assert, so we need to compute\n  # ``not pred`` here.\n  minus_one = _ir_constant(-1, mgpu_utils.dtype_to_ir_type(jnp.bool))\n  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:\n    pred = pred.registers.item()\n  not_pred = arith_dialect.xori(pred, minus_one)\n  cf_dialect.assert_(not_pred, exception.fmt_string)\n  return []\n\n\n@register_lowering_rule(pjit.relayout_p, mgpu.LoweringSemantics.Lane)\ndef _relayout_lowering_lane(\n    ctx: LoweringRuleContext, x, *, dst_layout\n):\n  del ctx, x, dst_layout\n  raise NotImplementedError(\n      \"relayout_p is not supported with Lane semantics.\"\n  )\n\n\n@register_lowering_rule(pjit.relayout_p, mgpu.LoweringSemantics.Warpgroup)\ndef _relayout_lowering_wg(\n    ctx: LoweringRuleContext, x, *, dst_layout\n):\n  if dst_layout is jax_layout.AutoLayout:\n    return x\n  layout = fa.TiledLayout(\n      dst_layout.tiling,\n      dst_layout.warp_dims,\n      dst_layout.lane_dims,\n      dst_layout.vector_dim,\n  )\n  if ctx.avals_in[0].ndim == 0:  # scalar case\n    if layout != mgpu.WGSplatFragLayout():","sourceCodeStart":4617,"sourceCodeEnd":4653,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L4617-L4653","documentation":"relayout_p (an automatic layout conversion primitive) has no lowering under Lane semantics in the Mosaic GPU backend. Lane-level code must have explicit layouts; the compiler refuses implicit relayout at that granularity.","triggerScenarios":"A computation running at Lane LoweringSemantics contains a relayout_p primitive, typically inserted automatically when layouts of producer/consumer mismatch inside warp-specialized or manually lowered code.","commonSituations":"Writing custom lowering rules or manually building FragmentedArrays with mismatched layouts; mixing layouts across warp-specialized branches.","solutions":["Insert an explicit layout_cast at the needed point instead of relying on relayout","Align layouts of the producer and consumer so relayout_p is never inserted","Move the operation to Warpgroup semantics where relayout is supported"],"exampleFix":"# before\ny = some_lane_level_op(x)  # triggers implicit relayout_p\n# after\nx2 = plgpu.layout_cast(x, target_layout)\ny = some_lane_level_op(x2)","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":["Insert explicit layout_cast instead of relying on relayout in lane-level code","Keep producer/consumer layouts consistent"],"tags":["pallas","mosaic-gpu","layout","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}