jax-ml/jax · error · ValueError

Expected the same number of out_{attr_suffix} ({len(out_layo

Error message

Expected the same number of out_{attr_suffix} ({len(out_layouts)}) as {value_type} results ({num_matching_results}). op=
  {op}

What it means

Same invariant check as for inputs, but for results: the number of out_layouts attributes must equal the number of vector-typed results of the op. A mismatch means the op's layout attributes were not built consistently with its results.

Source

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

    op: ir.OpView,
) -> None:
  """Ensures that the right number of in/out layouts are provided for an op.

  Layouts here are can be vector layouts, TMEM layouts, or SMEM transforms.
  """
  layouts = lambda attr: op.attributes[attr] if attr in op.attributes else []
  in_layouts = layouts(f"in_{attr_suffix}")
  out_layouts = layouts(f"out_{attr_suffix}")

  num_matching_operands = sum(map(filter_fn, op.operands))
  if len(in_layouts) != num_matching_operands:
    raise ValueError(
        f"Expected the same number of in_{attr_suffix} ({len(in_layouts)}) as "
        f"{value_type} operands ({num_matching_operands}). op=\n  {op}"
    )
  num_matching_results = sum(map(filter_fn, op.results))
  if len(out_layouts) != num_matching_results:
    raise ValueError(
        f"Expected the same number of out_{attr_suffix} ({len(out_layouts)}) "
        f"as {value_type} results ({num_matching_results}). op=\n  {op}"
    )


@dataclasses.dataclass(frozen=True)
class _TypeAndLayout:
  type: ir.Type
  layout: cs.Constant


def assign_layouts(solution: dict[ValueSite, cs.Constant]) -> None:
  """Assigns the layouts in `solution` to the MLIR ops they belong to.

  This function requires that, for each MLIR op that appears in `solution`,
  `solution` contains a layout assignment for all of its `vector`, TMEM, and
  SMEM operands and results. Block arguments are ignored.
  """

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use the DSL wrappers (e.g. mgpu helpers or FraggedArray-returning APIs) that infer out layouts
Fix the custom op builder so one out layout is emitted per vector result
Report/check upstream if it reproduces with library-provided ops

Defensive patterns

Strategy: validation

Validate before calling

n_res = sum(1 for r in op.results if isinstance(r.type, ir.VectorType))
assert len(op.attributes.get('out_layouts', [])) == n_res

Prevention

Emit one out_layout per vector result in custom lowering rules
Run the module through Mosaic's verifier early in development

When it happens

Trigger: Manually crafting an mgpu op whose out_layouts array length differs from its vector results count; ops produced by buggy lowering rules.

Common situations: Writing custom Mosaic lowering/layout rules, or hitting a regression after upgrading jax where result layout attr emission changed.

Related errors

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