jax-ml/jax · error · ValueError

Array layout mismatch: expected {v.layout} got {ty.layout.to

Error message

Array layout mismatch: expected {v.layout} got {ty.layout.to_mgpu()}.

What it means

Under lane semantics, a FragmentedArray matched against a ShapeDtypeStruct must carry exactly the layout obtained from the declared type's layout.to_mgpu(); a different register layout triggers this error.

Source

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

  raise NotImplementedError("inline_mgpu_p does not support discharge.")


def _type_check_mgpu_lane_semantics(v, ty):
  match (ty, v):
    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]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the inline function emit the array with the declared layout (apply mgpu layout conversion / relayout)
2. Update the declared Layout in ShapeDtypeStruct to the layout actually produced

Example fix

# before
return_type=ShapeDtypeStruct(s, dt, Layout.WGM_ROW)
# after
return_type=ShapeDtypeStruct(s, dt, actual_layout)  # or relayout the array in f

assert v.layout == declared.layout.to_mgpu()

When it happens

Trigger: Declaring a ShapeDtypeStruct with a Layout (e.g. row-major) while the inline function produces a FragmentedArray in a sliced/col-major/different layout.

Common situations: Mixing layouts after transposes or slicing inside the inline function without relayout.

Related errors

• Unexpected layout for {v} (expected: {ty})
• 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/d376ee7d93e6cd17. Report an issue: GitHub.