jax-ml/jax · error · ValueError

Inferred layout not found for operand {operand}.

Error message

Inferred layout not found for operand {operand}.

What it means

After solving the layout constraint system, every vector operand of an op with explicitly set in_layouts must have an inferred layout recorded in layout_for_variable. If a variable never got a layout assigned (e.g. it bypassed inference), this internal-consistency error fires.

Source

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

  return layout1 != layout2 and not isinstance(layout1, fa.WGSplatFragLayout)


def check_for_expensive_relayout(module: ir.Module):
  """Returns whether the given module has an expensive relayout."""
  layout_for_variable: dict[ir.Value, fa.FragmentedLayout] = {}

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure all values flowing into ops with explicit layouts are produced by ops that participate in layout inference
2. Rebuild the computation purely within the Mosaic DSL rather than mixing raw MLIR values
3. Update jax — missing-inferred-layout cases for constants are bug-fixed over time

try:
    run_layout_inference(module)
except ValueError as e:
    if 'Inferred layout not found' in str(e):
        # restructure kernel to avoid raw MLIR values feeding layout ops
        raise

When it happens

Trigger: An op with in_layouts set consumes a vector value that was never registered in the constraint system (e.g. a constant or op excluded from layout inference), so lookup by ValueSite fails.

Common situations: Custom ops or manually inserted MLIR values feeding mgpu ops; edge cases with constants/splat materialization in Mosaic kernels.

Related errors

• Unsupported reduction kind: {reduce}
• scan number of arguments doesn't match the number of jaxpr a
• Unrecognized index type: {typ}
• static_slice: unrecognized index {pidx.index} at position {p
• static_slice: unrecognized index {pidx.index}

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