{"record":{"id":"80764678cbc48d80","repo":"jax-ml/jax","slug":"array-dtype-mismatch-expected-v-mlir-dtype-got","errorCode":null,"errorMessage":"Array dtype mismatch: expected {v.mlir_dtype} got {mlir_dtype}.","messagePattern":"Array dtype mismatch: expected (.+?) got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3529,"sourceCode":"          if isinstance(r, state.AbstractRef)\n      ),\n  }\n\n\n@discharge.register_discharge_rule(inline_mgpu_p)\ndef _inline_mgpu_discharge(*args, **kwargs):\n  del args, kwargs\n  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(","sourceCodeStart":3511,"sourceCodeEnd":3547,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3511-L3547","documentation":"During inline_mgpu lowering under lane (non-warp-group) semantics, each FragmentedArray value is type-checked against its declared ShapeDtypeStruct; the element type of the runtime MLIR value must equal the declared JAX dtype converted to an MLIR type, otherwise this dtype mismatch error is raised.","triggerScenarios":"Declaring return_type/arg ShapeDtypeStruct with dtype jnp.float32 but the wrapped function produces an f64 or bf16 FragmentedArray.","commonSituations":"Implicit dtype promotion inside the inline function; declaring bf16 vs f16 incorrectly.","solutions":["Align the declared dtype with what the MLIR op produces (cast inside the function or fix the annotation)","Insert an explicit cast (arith.extf/trunc or mgpu cast) inside the wrapped function"],"exampleFix":"# before\ninline_mgpu(f, arg_types=[...], return_type=ShapeDtypeStruct(s, jnp.float32))\n# after\ninline_mgpu(f, arg_types=[...], return_type=ShapeDtypeStruct(s, jnp.bfloat16))  # match actual dtype","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nfrom jax._src.pallas.mosaic_gpu import mgpu_utils\nassert v.mlir_dtype == mgpu_utils.dtype_to_ir_type(declared.dtype)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Cast explicitly inside inline functions","Pin dtypes in ShapeDtypeStruct to the MLIR op output type"],"tags":["jax","pallas","inline-mgpu","dtype-mismatch"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}