jax-ml/jax · error · ValueError

Unexpected layout for {v} (expected: {ty})

Error message

Unexpected layout for {v} (expected: {ty})

What it means

When an argument/return type is given as a bare layout (SomeLayout) rather than a full ShapeDtypeStruct, the runtime FragmentedArray's layout must exactly equal that layout's mgpu equivalent; otherwise this 'Unexpected layout' error fires.

Source

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

    case (RefType(), ir.Value()) if isinstance(v.type, ir.MemRefType):
      pass
    case (ShapeDtypeStruct(), mgpu.FragmentedArray()):
      mlir_dtype = mgpu_utils.dtype_to_ir_type(ty.dtype)
      if v.mlir_dtype != mlir_dtype:
        raise ValueError(
            f"Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}."
        )
      if ty.shape != v.shape:
        raise ValueError(
            f"Array shape mismatch: expected {ty.shape} got {v.shape}."
        )
      if v.layout != ty.layout.to_mgpu():
        raise ValueError(
            f"Array layout mismatch: expected {v.layout} got {ty.layout.to_mgpu()}."
        )
    case (SomeLayout(), mgpu.FragmentedArray()):
      if ty.to_mgpu() != v.layout:
        raise ValueError(f"Unexpected layout for {v} (expected: {ty})")
    case _:
      raise ValueError(f"Unexpected type {ty} for value {v}")


def _inline_mgpu_flat_transformed_args(
    ctx: lowering.LoweringRuleContext,
    flat_args_and_transforms,
    flat_arg_types,
    pytree_args,
    pytree_ref_transforms,
  ) -> Sequence[ir.Value | mgpu.FragmentedArray]:
  flat_args = flat_args_and_transforms[:pytree_args.num_leaves]
  flat_arg_avals = ctx.avals_in[:pytree_args.num_leaves]
  ref_transforms = pytree_ref_transforms.unflatten(flat_args_and_transforms[pytree_args.num_leaves:])
  ref_transform_avals = pytree_ref_transforms.unflatten(ctx.avals_in[pytree_args.num_leaves:])
  is_wg_semantics = (
      ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the value with the exact declared layout (apply relayout before the call)
2. Relax the declared layout to the one actually produced

Example fix

# before
f = inline_mgpu(impl, arg_types=[some_layout], ...)
f(mgpu.splat(x, other_layout))
# after
f(mgpu.splat(x, some_layout.to_mgpu()))

assert ty.to_mgpu() == v.layout

When it happens

Trigger: Declaring arg_types=[plgpu.Layout(...)] and passing a FragmentedArray whose register layout differs (e.g. after a transpose or different vectorization).

Common situations: Splatting scalars with a default layout that does not match the declared one; layout drift after arithmetic that changes the layout.

Related errors

• Array layout mismatch: expected {v.layout} got {ty.layout.to
• inline_mgpu_p only supports plgpu.ShapeDtypeStruct return ty
• inline_mgpu_p only supports only SomeLayout and RefType arg
• Mismatched type shape: {treedef} != {treedef_ty}
• Mismatched type: {a, t}

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