jax-ml/jax · error · ValueError

Expected the same number of in_{attr_suffix} ({len(in_layout

Error message

Expected the same number of in_{attr_suffix} ({len(in_layouts)}) as {value_type} operands ({num_matching_operands}). op=
  {op}

What it means

During layout inference finalization, Mosaic checks that the number of in_layouts/in_vectors/in_indices attributes on an op equals the number of vector-typed operands. A mismatch means the op was constructed with missing or extra layout attributes, which is an invariant violation of the DSL.

Source

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


def _ensure_right_number_of_layouts(
    filter_fn: Callable[[ir.Value], bool],
    attr_suffix: str,
    value_type: str,
    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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Construct the op through the public mgpu Python helpers instead of raw MLIR op creation so layouts are attached automatically
Check for duplicate/missing entries in the in_* attribute array and make it match the count of vector operands
Update JAX/Mosaic to the latest version in case the op signature changed

Defensive patterns

Strategy: validation

Validate before calling

n_vec = sum(1 for o in op.operands if isinstance(o.type, ir.VectorType))
assert len(op.attributes.get('in_layouts', [])) == n_vec

Prevention

Prefer DSL helpers over raw MLIR op construction
Add assertions on attr-count vs operand-count in custom op builders

When it happens

Trigger: Building an mgpu dialect op manually (via ir insertion or by bypassing the Python DSL wrappers) with fewer/more in_layouts entries than vector operands; or a bug in Mosaic's own op construction for a new op type.

Common situations: Extending Mosaic with new custom ops, mixing MLIR-level op building with the Python DSL, or version skew after layout attr naming changes (in_layouts vs in_vectors).

Related errors

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