jax-ml/jax · error · NotImplementedError

Unsupported operand type: {element_type}

Error message

Unsupported operand type: {element_type}

What it means

Mosaic's mma supports only bf16, f16, f8e4m3fn, f8e5m2, i8 and i4 operand types. Any other element type raises NotImplementedError.

Source

Thrown at jax/experimental/mosaic/gpu/mma.py:211

  if n != n2:
    raise ValueError(f"N mismatch: {n} != {n2}")
  if k != k2:
    raise ValueError(f"K mismatch: {k} != {k2}")

  # todo(cperivol): A tile shape can have dimensions that are higher
  # multiples of the mma op size as long as those dimensions are not
  # sharded across warps.
  i4 = ir.IntegerType.get_signless(4)
  i8 = ir.IntegerType.get_signless(8)
  i32 = ir.IntegerType.get_signless(32)
  bf16 = ir.BF16Type.get()
  f16 = ir.F16Type.get()
  f8e4m3fn = ir.Float8E4M3FNType.get()
  f8e5m2 = ir.Float8E5M2Type.get()
  if (element_type := a.mlir_dtype) != b.mlir_dtype:
    raise ValueError(f"Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}")
  if element_type not in (bf16, f16, f8e4m3fn, f8e5m2, i8, i4):
    raise NotImplementedError(f"Unsupported operand type: {element_type}")
  if isinstance(element_type, ir.IntegerType):
    if acc.mlir_dtype != i32:
      raise NotImplementedError("Only s32 accumulator supported for integer operands.")
    if not acc.is_signed:
      raise ValueError("Only signed accumulator supported for integer operands.")
  elif acc.mlir_dtype != ir.F32Type.get():
    raise NotImplementedError("Only f32 accumulator supported for floating operands.")

  can_infer_from_acc_layout = (
      isinstance(acc.layout, fa.TiledLayout)
      and len(acc.layout.base_tile_shape) == 2
      and acc.layout.base_tile_shape[0] % 16 == 0
  )
  if not can_infer_from_acc_layout:
    raise ValueError("Expected MMALayouts.acc for acc")
  m_warps = acc.layout.base_tile_shape[0] // 16  # type: ignore
  layouts = MMALayouts(element_type, m_warps=m_warps)
  if layouts.lhs != a.layout:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Downcast operands to a supported dtype (bf16/f16/f8/int8) before mma
2. Use plain elementwise multiply-accumulate for unsupported high-precision types
3. Keep an f32 accumulator with f16/bf16 operands for range

Example fix

// before
acc = mma.mma(a_f32, b_f32, acc)
// after
acc = mma.mma(a_f32.astype(jnp.bfloat16), b_f32.astype(jnp.bfloat16), acc)

supported = {'bf16','f16','f8e4m3fn','f8e5m2','i8','i4'}
assert str(a.mlir_dtype) in supported

When it happens

Trigger: Calling mma() with f64, f32, or an exotic integer width (e.g. i16) as operand element type.

Common situations: Porting kernels that assume tf32/f32 tensor cores; older GPUs or expecting full float32 MMA support.

Related errors

• m_warps must be 1, 2, or 4, but got {m_warps=}
• is_signed must be specified for integer types
• M mismatch: {m} != {m2}
• N mismatch: {n} != {n2}
• K mismatch: {k} != {k2}

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