{"record":{"id":"2e0c6b673a9150bf","repo":"jax-ml/jax","slug":"mma-with-element-type-elem-type-str-only-support","errorCode":null,"errorMessage":"MMA with element type {elem_type_str} only supports accumulators of type f32, but got: {d.dtype}","messagePattern":"MMA with element type (.+?) only supports accumulators of type f32, but got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":298,"sourceCode":"  else:\n    raise ValueError(f\"Only M=128 and M=64 are supported for MMA, but got M={m}\")\n  f32 = ir.F32Type.get()\n  f16 = ir.F16Type.get()\n  s32 = ir.IntegerType.get_signless(32)\n  elem_type_str = (\n      f\"{a_element_type}\"\n      if a_element_type == b_element_type\n      else f\"({a_element_type}, {b_element_type})\"\n  )\n  if a_element_type == f32 or a_element_type == ir.BF16Type.get():\n    if a_element_type == f32 and is_sparse:\n      raise NotImplementedError(\"Sparse MMA unsupported for f32\")\n    if is_scaled:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} does not support block scaling\"\n      )\n    if d.dtype != f32:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} only supports accumulators\"\n          f\" of type f32, but got: {d.dtype}\"\n      )\n  elif a_element_type == f16:\n    if is_scaled:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} does not support block scaling\"\n      )\n    if d.dtype != f16 and d.dtype != f32:\n      raise ValueError(\n          f\"MMA with element type {elem_type_str} only supports accumulators of\"\n          f\" type f32 or f16, but got: {d.dtype}\"\n      )\n  elif any(\n      isinstance(a_element_type, t)\n      for t in {ir.Float8E5M2Type, ir.Float8E4M3FNType}\n  ):\n    if d.dtype != f16 and d.dtype != f32:","sourceCodeStart":280,"sourceCodeEnd":316,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L280-L316","documentation":"When the A operand is f32 or bf16, the tcgen05 MMA instruction accumulates in f32 only. If the provided accumulator tensor d has any other dtype (e.g. f16), the op raises this ValueError before emitting MLIR.","triggerScenarios":"Calling tcgen05.mma with a/b of type f32 or bf16 and an accumulator d whose dtype is not f32 (commonly f16).","commonSituations":"Reusing an f16 accumulator allocation from an fp16 kernel in a bf16 kernel; allocating TMEM with a default dtype that doesn't match the operand type; mixed-precision experiments with f16 accumulation.","solutions":["Allocate/convert the accumulator d to f32 (jnp.float32 / ir.F32Type)","Match d's dtype to f32 whenever A is f32 or bf16","If f16 accumulation is required, switch the operands themselves to f16 instead"],"exampleFix":"# before\nd = tmem.alloc((m, n), dtype=jnp.float16)\ntcgen05.mma(a_bf16, b_bf16, d)  # raises\n# after\nd = tmem.alloc((m, n), dtype=jnp.float32)\ntcgen05.mma(a_bf16, b_bf16, d)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nassert d.dtype == jnp.float32, f'f32/bf16 operands require f32 accumulator, got {d.dtype}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive accumulator dtype from operand dtype with a lookup helper","Avoid default dtypes in tmem.alloc; always pass dtype explicitly"],"tags":["jax","mosaic","gpu","accumulator-dtype","dtype","mma","tcgen05"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}