{"record":{"id":"daaf190862a886eb","repo":"jax-ml/jax","slug":"dtype-mismatch-a-mlir-dtype-b-mlir-dtype","errorCode":null,"errorMessage":"Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}","messagePattern":"Dtype mismatch: (.+?) != (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/mma.py","lineNumber":209,"sourceCode":"  if m != m2:\n    raise ValueError(f\"M mismatch: {m} != {m2}\")\n  if n != n2:\n    raise ValueError(f\"N mismatch: {n} != {n2}\")\n  if k != k2:\n    raise ValueError(f\"K mismatch: {k} != {k2}\")\n\n  # todo(cperivol): A tile shape can have dimensions that are higher\n  # multiples of the mma op size as long as those dimensions are not\n  # sharded across warps.\n  i4 = ir.IntegerType.get_signless(4)\n  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  bf16 = ir.BF16Type.get()\n  f16 = ir.F16Type.get()\n  f8e4m3fn = ir.Float8E4M3FNType.get()\n  f8e5m2 = ir.Float8E5M2Type.get()\n  if (element_type := a.mlir_dtype) != b.mlir_dtype:\n    raise ValueError(f\"Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}\")\n  if element_type not in (bf16, f16, f8e4m3fn, f8e5m2, i8, i4):\n    raise NotImplementedError(f\"Unsupported operand type: {element_type}\")\n  if isinstance(element_type, ir.IntegerType):\n    if acc.mlir_dtype != i32:\n      raise NotImplementedError(\"Only s32 accumulator supported for integer operands.\")\n    if not acc.is_signed:\n      raise ValueError(\"Only signed accumulator supported for integer operands.\")\n  elif acc.mlir_dtype != ir.F32Type.get():\n    raise NotImplementedError(\"Only f32 accumulator supported for floating operands.\")\n\n  can_infer_from_acc_layout = (\n      isinstance(acc.layout, fa.TiledLayout)\n      and len(acc.layout.base_tile_shape) == 2\n      and acc.layout.base_tile_shape[0] % 16 == 0\n  )\n  if not can_infer_from_acc_layout:\n    raise ValueError(\"Expected MMALayouts.acc for acc\")\n  m_warps = acc.layout.base_tile_shape[0] // 16  # type: ignore","sourceCodeStart":191,"sourceCodeEnd":227,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/mma.py#L191-L227","documentation":"The two MMA operands must have identical MLIR dtypes; a.mlir_dtype != b.mlir_dtype raises ValueError. Mixed-precision matmul at this low level is not supported.","triggerScenarios":"Calling mma() with a as bf16 and b as fp16, or one operand i8 and the other i4.","commonSituations":"Converting only one operand to half precision, or loading one operand from a differently-typed tensor (e.g. via constants) in a Mosaic kernel.","solutions":["Cast both operands to the same dtype (e.g. both bf16) before mma","Ensure tensor→FraggedArray conversion preserves dtype for both operands","Check for implicit dtype promotion in your data pipeline"],"exampleFix":"// before\nacc = mma.mma(a_bf16, b_f16, acc)\n// after\nb_bf16 = (b_f16.astype(jnp.bfloat16))\nacc = mma.mma(a_bf16, b_bf16, acc)","handlingStrategy":"validation","validationCode":"assert a.mlir_dtype == b.mlir_dtype, 'mma operands need same dtype'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Standardize both operands to one dtype before building FragmentedArrays"],"tags":["jax","mosaic","mma","dtype-mismatch"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}