jax-ml/jax · error · ValueError

M mismatch: {m} != {m2}

Error message

M mismatch: {m} != {m2}

What it means

The low-level mma helper requires the M dimension of the accumulator to match the M of operand a (a.shape[0] == acc.shape[0]). A mismatch means the shapes don't form a valid matmul tile.

Source

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

  Args:
    acc: A `FragmentedArray` with a `TiledLayout` generated from
      `MMALayouts.acc`.
    a: A `FragmentedArray` with a `TiledLayout`  generated from
      `MMALayouts.lhs`.
    b: A `FragmentedArray` with a `TiledLayout` generated from `MMALayouts.rhs`.

  Returns:
    A new `FragmentedArray` with the result of the computation with
      the same type as `acc`.
  """

  (m, k) = a.shape
  (k2, n) = b.shape
  (m2, n2) = acc.shape

  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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate acc with shape (a.shape[0], b.shape[1])
2. Fix the operand tiling so M matches the accumulator
3. Zero-init a correctly shaped accumulator each matmul

Example fix

// before
acc = fa.from_tensor(jnp.zeros((32, n), jnp.float32))
acc = mma.mma(a_64xk, b, acc)
// after
acc = fa.from_tensor(jnp.zeros((64, n), jnp.float32))
acc = mma.mma(a_64xk, b, acc)

assert a.shape[0] == acc.shape[0], 'M mismatch'

When it happens

Trigger: Calling mma(a, b, acc) where acc has fewer/more rows than a, e.g. acc of shape (32, N) with a of shape (64, K).

Common situations: Accumulator allocated with a different tile shape than the operands, or reusing an accumulator across matmuls of different M.

Related errors

• N mismatch: {n} != {n2}
• K mismatch: {k} != {k2}
• Logical shape {self.logical_shape} cannot be bigger than con
• Accumulator and LHS have incompatible shapes. Expected LHS t
• Accumulator and RHS have incompatible shapes. Expected RHS t

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