jax-ml/jax · error · ValueError

N mismatch: {n} != {n2}

Error message

N mismatch: {n} != {n2}

What it means

The N dimension of the accumulator must equal b.shape[1] (acc.shape[1] == b.shape[1]). A mismatch indicates the matmul tile's output width doesn't line up with operand b.

Source

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

    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):
    raise NotImplementedError(f"Unsupported operand type: {element_type}")
  if isinstance(element_type, ir.IntegerType):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate/reallocate acc with shape (m, b.shape[1])
2. Keep N of b and acc consistent when retuning tile sizes

Example fix

// before
acc = fa.from_tensor(jnp.zeros((m, 64), jnp.float32))
acc = mma.mma(a, b_n128, acc)
// after
acc = fa.from_tensor(jnp.zeros((m, 128), jnp.float32))
acc = mma.mma(a, b_n128, acc)

assert b.shape[1] == acc.shape[1], 'N mismatch'

When it happens

Trigger: Calling mma(a, b, acc) with an accumulator whose second dim differs from b's second dim, e.g. N-tile changed without resizing acc.

Common situations: Sweeping N tile sizes while reusing a stale accumulator buffer, or transposing b without adjusting acc.

Related errors

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