jax-ml/jax · error · ValueError

K mismatch: {k} != {k2}

Error message

K mismatch: {k} != {k2}

What it means

The contraction dimension must agree: a.shape[1] must equal b.shape[0]. A K mismatch means the operand tiles cannot be multiplied.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Slice a and b with the same K extent in each mma call
2. Check loop bounds for the contraction dimension
3. Re-verify tile shapes after reshaping operands

Example fix

// before
acc = mma.mma(a[:, :32], b[64:, :], acc)
// after
acc = mma.mma(a[:, :32], b[:32, :], acc)

assert a.shape[1] == b.shape[0], 'K mismatch'

When it happens

Trigger: Calling mma(a, b, acc) where a's K dim differs from b's K dim, e.g. slicing one operand differently along the contraction axis.

Common situations: Split-K or partial-accumulation loops slicing a and b inconsistently; off-by-one loop bounds over K.

Related errors

• M mismatch: {m} != {m2}
• N mismatch: {n} != {n2}
• 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/ea4fc60343044759. Report an issue: GitHub.