jax-ml/jax · error · NotImplementedError

relayout_p is not supported with Lane semantics.

Error message

relayout_p is not supported with Lane semantics.

What it means

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.

Source

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

  exception = jax.tree.unflatten(exception_tree, payload)
  assert isinstance(exception, checkify.FailedCheckError)

  # check_p has an inverted predicate compared to assert, so we need to compute
  # ``not pred`` here.
  minus_one = _ir_constant(-1, mgpu_utils.dtype_to_ir_type(jnp.bool))
  if ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Lane:
    pred = pred.registers.item()
  not_pred = arith_dialect.xori(pred, minus_one)
  cf_dialect.assert_(not_pred, exception.fmt_string)
  return []


@register_lowering_rule(pjit.relayout_p, mgpu.LoweringSemantics.Lane)
def _relayout_lowering_lane(
    ctx: LoweringRuleContext, x, *, dst_layout
):
  del ctx, x, dst_layout
  raise NotImplementedError(
      "relayout_p is not supported with Lane semantics."
  )


@register_lowering_rule(pjit.relayout_p, mgpu.LoweringSemantics.Warpgroup)
def _relayout_lowering_wg(
    ctx: LoweringRuleContext, x, *, dst_layout
):
  if dst_layout is jax_layout.AutoLayout:
    return x
  layout = fa.TiledLayout(
      dst_layout.tiling,
      dst_layout.warp_dims,
      dst_layout.lane_dims,
      dst_layout.vector_dim,
  )
  if ctx.avals_in[0].ndim == 0:  # scalar case
    if layout != mgpu.WGSplatFragLayout():

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Insert an explicit layout_cast at the needed point instead of relying on relayout
2. Align layouts of the producer and consumer so relayout_p is never inserted
3. Move the operation to Warpgroup semantics where relayout is supported

Example fix

# before
y = some_lane_level_op(x)  # triggers implicit relayout_p
# after
x2 = plgpu.layout_cast(x, target_layout)
y = some_lane_level_op(x2)

When it happens

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

Common situations: Writing custom lowering rules or manually building FragmentedArrays with mismatched layouts; mixing layouts across warp-specialized branches.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• packed cannot be specified if layout is specified.
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Unsupported transform: {type(transform)}

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