{"record":{"id":"d376ee7d93e6cd17","repo":"jax-ml/jax","slug":"array-layout-mismatch-expected-v-layout-got-ty","errorCode":null,"errorMessage":"Array layout mismatch: expected {v.layout} got {ty.layout.to_mgpu()}.","messagePattern":"Array layout mismatch: expected (.+?) got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3537,"sourceCode":"  raise NotImplementedError(\"inline_mgpu_p does not support discharge.\")\n\n\ndef _type_check_mgpu_lane_semantics(v, ty):\n  match (ty, v):\n    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]","sourceCodeStart":3519,"sourceCodeEnd":3555,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3519-L3555","documentation":"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.","triggerScenarios":"Declaring a ShapeDtypeStruct with a Layout (e.g. row-major) while the inline function produces a FragmentedArray in a sliced/col-major/different layout.","commonSituations":"Mixing layouts after transposes or slicing inside the inline function without relayout.","solutions":["Make the inline function emit the array with the declared layout (apply mgpu layout conversion / relayout)","Update the declared Layout in ShapeDtypeStruct to the layout actually produced"],"exampleFix":"# before\nreturn_type=ShapeDtypeStruct(s, dt, Layout.WGM_ROW)\n# after\nreturn_type=ShapeDtypeStruct(s, dt, actual_layout)  # or relayout the array in f","handlingStrategy":"validation","validationCode":"assert v.layout == declared.layout.to_mgpu()","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Relayout FragmentedArrays to the declared layout before return","Keep layout round-trips through to_mgpu consistent"],"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"}