jax-ml/jax · error · ValueError

Incompatible shapes for matrix multiplication: lhs={a.shape}

Error message

Incompatible shapes for matrix multiplication: lhs={a.shape}, rhs={b.shape=}, acc={acc.shape}

What it means

wgmma validates that the accumulator shape (m, n) and lhs (m, k), rhs (k, n) form a consistent matrix multiplication before emitting the instruction. Mismatched dimensions are a pure user API error.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:1852

  Conceptually, this is equivalent to doing ``acc[...] += a[...] @ b[...]``,
  except that the computation is performed asynchronously.

  Args:
    acc: The accumulator reference. Needs to be allocated via
      :func:`jax.experimental.pallas.run_scoped` called with a
      :func:`jax.experimental.pallas.mosaic_gpu.WGMMAAccumulatorRef`.
    a: The left hand side operand reference.
    b: The right hand side operand reference.

  See also:
    :func:`jax.experimental.pallas.mosaic_gpu.wgmma_wait`
  """
  m, n = acc.shape
  m2, k = a.shape
  k2, n2 = b.shape

  if m != m2 or n != n2 or k != k2:
    raise ValueError(
        f"Incompatible shapes for matrix multiplication: lhs={a.shape},"
        f" rhs={b.shape=}, acc={acc.shape}"
    )

  # A and B must share a dtype, except that the e4m3/e5m2 FP8 pair may be mixed:
  # `wgmma` takes independent `.atype`/`.btype` operands for FP8.
  fp8_dtypes = (jnp.float8_e4m3fn, jnp.float8_e5m2)
  both_fp8 = a.dtype in fp8_dtypes and b.dtype in fp8_dtypes
  if a.dtype != b.dtype and not both_fp8:
    raise ValueError(
        "Mixed input dtypes for matrix multiplication unsupported: "
        f"lhs={a.dtype}, rhs={b.dtype}"
    )

  acc_transforms_leaves: list
  if isinstance(acc, pallas_core.TransformedRef):
    acc_transforms_leaves, acc_transforms_tree = jax.tree.flatten(acc.transforms)
    acc = acc.ref

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Check shapes: allocate acc as jnp.zeros((a.shape[0], b.shape[1]), ...) and ensure a.shape[1] == b.shape[0].
Verify operand order — lhs must be (m, k), rhs (k, n).
Add an assert in the kernel to fail fast with clearer context.

Example fix

# before
acc = smem.zeros((n_dim, m_dim), jnp.float32)
out = wgmma(a, b, acc)
# after
acc = smem.zeros((a.shape[0], b.shape[1]), jnp.float32)
out = wgmma(a, b, acc)

Defensive patterns

Strategy: validation

Validate before calling

m, k = a.shape; k2, n = b.shape
assert k == k2, f'inner dims {k} vs {k2}'
assert acc.shape == (m, n), f'acc {acc.shape} != {(m, n)}'

Prevention

Allocate the accumulator from lhs/rhs shapes in the same expression as the wgmma call.
Add shape asserts in kernel setup code.

When it happens

Trigger: Calling wgmma(a, b, acc) where acc.shape != (a.shape[0], b.shape[1]) or a.shape[1] != b.shape[0].

Common situations: Swapped operands (b, a instead of a, b), wrong accumulator allocation after changing block sizes, or forgetting that wgmma expects (m,k)@(k,n).

Related errors

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