{"record":{"id":"6d7cb66dd75cd35f","repo":"jax-ml/jax","slug":"unexpected-layout-for-v-expected-ty","errorCode":null,"errorMessage":"Unexpected layout for {v} (expected: {ty})","messagePattern":"Unexpected layout for (.+?) \\(expected: (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3542,"sourceCode":"    case (RefType(), ir.Value()) if isinstance(v.type, ir.MemRefType):\n      pass\n    case (ShapeDtypeStruct(), mgpu.FragmentedArray()):\n      mlir_dtype = mgpu_utils.dtype_to_ir_type(ty.dtype)\n      if v.mlir_dtype != mlir_dtype:\n        raise ValueError(\n            f\"Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}.\"\n        )\n      if ty.shape != v.shape:\n        raise ValueError(\n            f\"Array shape mismatch: expected {ty.shape} got {v.shape}.\"\n        )\n      if v.layout != ty.layout.to_mgpu():\n        raise ValueError(\n            f\"Array layout mismatch: expected {v.layout} got {ty.layout.to_mgpu()}.\"\n        )\n    case (SomeLayout(), mgpu.FragmentedArray()):\n      if ty.to_mgpu() != v.layout:\n        raise ValueError(f\"Unexpected layout for {v} (expected: {ty})\")\n    case _:\n      raise ValueError(f\"Unexpected type {ty} for value {v}\")\n\n\ndef _inline_mgpu_flat_transformed_args(\n    ctx: lowering.LoweringRuleContext,\n    flat_args_and_transforms,\n    flat_arg_types,\n    pytree_args,\n    pytree_ref_transforms,\n  ) -> Sequence[ir.Value | mgpu.FragmentedArray]:\n  flat_args = flat_args_and_transforms[:pytree_args.num_leaves]\n  flat_arg_avals = ctx.avals_in[:pytree_args.num_leaves]\n  ref_transforms = pytree_ref_transforms.unflatten(flat_args_and_transforms[pytree_args.num_leaves:])\n  ref_transform_avals = pytree_ref_transforms.unflatten(ctx.avals_in[pytree_args.num_leaves:])\n  is_wg_semantics = (\n      ctx.module_ctx.lowering_semantics == mgpu.LoweringSemantics.Warpgroup\n  )","sourceCodeStart":3524,"sourceCodeEnd":3560,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3524-L3560","documentation":"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.","triggerScenarios":"Declaring arg_types=[plgpu.Layout(...)] and passing a FragmentedArray whose register layout differs (e.g. after a transpose or different vectorization).","commonSituations":"Splatting scalars with a default layout that does not match the declared one; layout drift after arithmetic that changes the layout.","solutions":["Pass the value with the exact declared layout (apply relayout before the call)","Relax the declared layout to the one actually produced"],"exampleFix":"# before\nf = inline_mgpu(impl, arg_types=[some_layout], ...)\nf(mgpu.splat(x, other_layout))\n# after\nf(mgpu.splat(x, some_layout.to_mgpu()))","handlingStrategy":"validation","validationCode":"assert ty.to_mgpu() == v.layout","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Splat with the exact declared layout","Avoid layout-changing ops immediately before inline_mgpu calls"],"tags":["jax","pallas","inline-mgpu","layout-mismatch"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}