jax-ml/jax · error · ValueError

MMA with element type

Error message

MMA with element type {elem_type_str} only supports accumulators of type f32 or f16, but got: {d.dtype}

What it means

For f32/bf16 operand types the tcgen05 MMA instruction only defines f32 accumulators. It fires when d.dtype is neither f32 nor f16 for that element type combination.

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

Example fix

# before
d = tmem.alloc((m, n), dtype=jnp.bfloat16)
tcgen05.mma(a_f16, b_f16, d)  # raises
# after
d = tmem.alloc((m, n), dtype=jnp.float32)
tcgen05.mma(a_f16, b_f16, d)

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp
assert d.dtype in (jnp.float16, jnp.float32), f'f16 operands need f16/f32 accumulator, got {d.dtype}'

Prevention

Never use bf16 accumulators with tcgen05 MMA
Write one shared assert helper for accumulator/operand dtype pairs

When it happens

Trigger: Calling tcgen05.mma with f16 operands and an accumulator d typed as bf16, fp32-variant other than F32, or any non f16/f32 type.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/0bdbe265bc0d2f07. Report an issue: GitHub.

Appendix: source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:304

      f"{a_element_type}"
      if a_element_type == b_element_type
      else f"({a_element_type}, {b_element_type})"
  )
  if a_element_type == f32 or a_element_type == ir.BF16Type.get():
    if a_element_type == f32 and is_sparse:
      raise NotImplementedError("Sparse MMA unsupported for f32")
    if is_scaled:
      raise ValueError(
          f"MMA with element type {elem_type_str} does not support block scaling"
      )
    if d.dtype != f32:
      raise ValueError(
          f"MMA with element type {elem_type_str} only supports accumulators"
          f" of type f32, but got: {d.dtype}"
      )
  elif a_element_type == f16:
    if is_scaled:
      raise ValueError(
          f"MMA with element type {elem_type_str} does not support block scaling"
      )
    if d.dtype != f16 and d.dtype != f32:
      raise ValueError(
          f"MMA with element type {elem_type_str} only supports accumulators of"
          f" type f32 or f16, but got: {d.dtype}"
      )
  elif any(
      isinstance(a_element_type, t)
      for t in {ir.Float8E5M2Type, ir.Float8E4M3FNType}
  ):
    if d.dtype != f16 and d.dtype != f32:
      raise ValueError(
          f"MMA with element type {elem_type_str} only supports accumulators of"
          f" type f32 or f16, but got: {d.dtype}"
      )
    if is_scaled and d.dtype != f32:
      raise ValueError(

View on GitHub (pinned to 1e1c6a8fc0)