{"record":{"id":"0bdbe265bc0d2f07","repo":"jax-ml/jax","slug":"mma-with-element-type-elem-type-str-only-support-0bdbe2","errorCode":null,"errorMessage":"MMA with element type {elem_type_str} only supports accumulators of type f32 or f16, but got: {d.dtype}","messagePattern":"MMA with element type (.+?) only supports accumulators of type f32 or f16, but got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":304,"sourceCode":"      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:\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    if is_scaled and d.dtype != f32:\n      raise ValueError(","sourceCodeStart":286,"sourceCodeEnd":322,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L286-L322","documentation":"For f16 A operands, the tcgen05 MMA can accumulate in either f16 or f32. Any other accumulator dtype (e.g. bf16 or an integer type) is invalid and rejected here.","triggerScenarios":"Calling tcgen05.mma with f16 operands and an accumulator d typed as bf16, fp32-variant other than F32, or any non f16/f32 type.","commonSituations":"Using bf16 accumulators because they work elsewhere in the pipeline; allocating TMEM accumulators with a global default dtype; dtype typos when constructing MLIR types manually.","solutions":["Set the accumulator d dtype to f16 or f32","Prefer f32 accumulation unless you specifically want f16 for TMEM bandwidth reasons","Centralize accumulator allocation so dtype is derived from the operand type"],"exampleFix":"# before\nd = tmem.alloc((m, n), dtype=jnp.bfloat16)\ntcgen05.mma(a_f16, b_f16, d)  # raises\n# after\nd = tmem.alloc((m, n), dtype=jnp.float32)\ntcgen05.mma(a_f16, b_f16, d)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nassert d.dtype in (jnp.float16, jnp.float32), f'f16 operands need f16/f32 accumulator, got {d.dtype}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never use bf16 accumulators with tcgen05 MMA","Write one shared assert helper for accumulator/operand dtype pairs"],"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"}