jax-ml/jax · error · ValueError

Inferred layout {operand_layout} for operand {operand} does

Error message

Inferred layout {operand_layout} for operand {operand} does not match the layout in layout_for_variable {layout_for_variable[operand]}.

What it means

The layout inferred by the constraint system for an operand disagrees with the op's explicitly requested in_layout, and the relayout between them is expensive (not a cheap bitcast/reshape). Mosaic refuses to silently insert costly data movement.

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:2701

  def _check_for_expensive_relayout(op: ir.OpView):
    if not inference_utils.should_have_layout(op):
      return
    if inference_utils.has_in_layouts_set(op):
      in_layouts = iter(layouts_lib.from_layout_attr(l)
                        for l in cast(ir.ArrayAttr, op.attributes["in_layouts"]))
      for operand in op.operands:
        assert isinstance(operand, ir.Value)
        if not isinstance(operand.type, ir.VectorType):
          continue
        if operand not in layout_for_variable:
          raise ValueError(
              f"Inferred layout not found for operand {operand}."
          )
        operand_layout = next(in_layouts)
        # TODO(bchetioui): refine to figure out whether it's a cheap relayout.
        if _is_expensive_relayout(layout_for_variable[operand], operand_layout):
          raise ValueError(
              f"Inferred layout {operand_layout} for operand {operand} does "
              f"not match the layout in layout_for_variable "
              f"{layout_for_variable[operand]}."
          )
    if inference_utils.has_out_layouts_set(op):
      out_layouts = iter(layouts_lib.from_layout_attr(l)
                         for l in cast(ir.ArrayAttr, op.attributes["out_layouts"]))
      for result in op.results:
        assert isinstance(result, ir.Value)
        if not isinstance(result.type, ir.VectorType):
          continue
        assert result not in layout_for_variable
        layout_for_variable[result] = next(out_layouts)

  for op in module.body:
    traverse_op(op, _check_for_expensive_relayout)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Remove the explicit layout/cast and let inference pick layouts
Insert an explicit layout_cast where you accept the cost, instead of relying on implicit conversion
Restructure the kernel so producer and consumer agree on one layout (e.g. keep the MMA-compatible layout throughout)

Example fix

// before
acc = mgpu.layout_cast(acc, some_other_layout)
out = op_with_explicit_in_layout(acc)
// after
out = op_with_explicit_in_layout(acc)  # drop the cast; keep one layout

Defensive patterns

Strategy: validation

Validate before calling

if _is_expensive_relayout(current_layout, requested_layout):
    x = mgpu.layout_cast(x, requested_layout)  # make the move explicit

Prevention

Avoid pinning conflicting layouts on both producer and consumer
Chain incompatible layouts via explicit layout_cast calls

When it happens

Trigger: User sets a layout cast or in_layouts that conflicts with the layouts propagation naturally derives, e.g. feeding a value laid out for MMA into an op expecting a different non-compatible tiled layout.

Common situations: Explicit layout_cast / set layouts on ops in Mosaic kernels that fight the inferred layouts; mixing warps/registers layouts across op boundaries.

Related errors

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