jax-ml/jax · error · ValueError

Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}

Error message

Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}

What it means

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.

Source

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

  if m != m2:
    raise ValueError(f"M mismatch: {m} != {m2}")
  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast both operands to the same dtype (e.g. both bf16) before mma
2. Ensure tensor→FraggedArray conversion preserves dtype for both operands
3. Check for implicit dtype promotion in your data pipeline

Example fix

// before
acc = mma.mma(a_bf16, b_f16, acc)
// after
b_bf16 = (b_f16.astype(jnp.bfloat16))
acc = mma.mma(a_bf16, b_bf16, acc)

assert a.mlir_dtype == b.mlir_dtype, 'mma operands need same dtype'

When it happens

Trigger: Calling mma() with a as bf16 and b as fp16, or one operand i8 and the other i4.

Common situations: Converting only one operand to half precision, or loading one operand from a differently-typed tensor (e.g. via constants) in a Mosaic kernel.

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/daaf190862a886eb. Report an issue: GitHub.