jax-ml/jax · error · ValueError

Output layout {out_layout} does not match the layout of the

Error message

Output layout {out_layout} does not match the layout of the returned fragmented array {fa.layout}.

What it means

For vector-typed outputs of inline_mgpu, the layout attribute computed from the declared output layout must equal the layout of the returned FragmentedArray. Divergent layouts (e.g. declared row-major but fragment is column-major) raise this ValueError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:3838

    args = jax.tree.unflatten(pytree_args, fn_inputs)
    inner_ret = mgpu_fn(ctx.launch_ctx, *args)
    if inner_ret is None:
      inner_ret = []
    elif not isinstance(inner_ret, tuple) and not isinstance(inner_ret, list):
      inner_ret = [inner_ret]
    ir_ret = []
    for fa, result_ty, out_layout in zip(
        inner_ret, results_ty, out_layouts, strict=True
    ):
      if not isinstance(fa, mgpu.FragmentedArray):
        raise ValueError(f"Expected a FragmentedArray, but got: {fa}")
      if isinstance(result_ty, ir.VectorType):
        result_shape = ir.VectorType(result_ty).shape
        if fa.shape != tuple(result_shape):
          raise ValueError(f"Expected {result_shape} but got {fa.shape}")
        if out_layout != mgpu.layouts.to_layout_attr(fa.layout):
          raise ValueError(
              f"Output layout {out_layout} does not match the layout of the"
              f" returned fragmented array {fa.layout}."
          )
        ir_ret.append(
            mgpu.dialect_lowering.fragmented_array_to_ir(fa, result_ty)
        )
      else:  # scalar case.
        assert out_layout is None
        if fa.shape:
          raise ValueError(f"Expected 0D shape, but got {fa.shape}")
        if not isinstance(fa.layout, mgpu.WGSplatFragLayout):
          raise ValueError(f"Expected WGSplatFragLayout, but got {fa.layout}")
        value = fa.registers.item()
        ir_ret.append(value)

    mgpu.dialect.return_(ir_ret)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set the declared output layout to match the fragment layout produced inside the callback
2. Insert an explicit layout conversion in the callback so the final layout matches the declaration

assert mgpu.layouts.to_layout_attr(fa.layout) == out_layout

When it happens

Trigger: Returning a FragmentedArray whose fa.layout differs from the out_layout derived from the declared return type or layout hint.

Common situations: Operations inside the callback change the fragment layout (transpose, layout conversion) without a matching declared out_layout; stale layout hints.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.
• Cannot specify both scratch_shapes and scratch_types. Use sc
• Unsupported dtype: {ref.dtype}

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