jax-ml/jax · error · ValueError

Incompatible shapes for matrix multiplication: lhs=

Error message

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

What it means

Shape-consistency guard in the wgmma wrapper: the accumulator must be (m, n), the lhs operand (m, k), and the rhs operand (k, n). The error fires when any of these dimension pairs disagree, meaning the operands cannot form a valid blocked matrix multiply accumulating into the supplied acc ref.

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.

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