jax-ml/jax · error · ValueError

Accumulator and RHS have incompatible shapes. Expected RHS…

Error message

Accumulator and RHS have incompatible shapes. Expected RHS to have shape (k, n) and accumulator to have shape (m, n). Accumulator: {acc.shape}. RHS: {b.shape}.

What it means

Non-collective shape guard in tcgen05_mma: outside collective mode the rhs operand must be (k, n) aligned with an accumulator of (m, n). The error fires when the rhs dimensions (rhs_k/rhs_n) don't match the accumulator's k/n, so the plain MMA shape contract is violated.

Solutions

Match accumulator N to b.shape[1] exactly
If porting from collective mode, halve the accumulator N
Validate all three shapes (acc (m,n), a (m,k), b (k,n)) before the call

Example fix

# before
acc = allocate(TMEM, (m, 2 * n), jnp.float32)  # copied from collective kernel
tcgen05.mma(a, b, acc, k_dim=k)
# after
acc = allocate(TMEM, (m, n), jnp.float32)
tcgen05.mma(a, b, acc, k_dim=k)

Defensive patterns

Strategy: validation

Validate before calling

assert acc.shape[1] == b.shape[1], (acc.shape, b.shape)
tcgen05.mma(a, b, acc, k_dim=a.shape[1])

Prevention

Add a shape-check helper used by all tcgen05.mma call sites
Avoid copying sizes between collective and non-collective kernels

When it happens

Trigger: Calling tcgen05.mma with accumulator (m, 128) but RHS block of shape (k, 64), or vice versa.

Common situations: Changing the N tiling of the B matrix without updating the accumulator allocation; copying accumulator sizing from a collective-mode kernel into a non-collective one (off-by-2x N).

Related errors

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

Appendix: source

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

  rhs_k, rhs_n = b.shape
  is_sparse = a_sparse_metadata is not None

  if acc_m != lhs_m:
    raise ValueError(
        "Accumulator and LHS have incompatible shapes. Expected LHS to have"
        " shape (m, k) and accumulator to have shape (m, n). Accumulator:"
        f" {acc.shape}. LHS: {a.shape}."
    )

  if collective_axis is not None:
    if acc_n != rhs_n * 2:
      raise ValueError(
          "Accumulator and RHS have incompatible shapes. Expected RHS to have "
          "shape (k, n) and accumulator to have shape (m, n * 2) in "
          f"collective mode. Accumulator: {acc.shape}. RHS: {b.shape}."
      )
  elif acc_n != rhs_n:
    raise ValueError(
        "Accumulator and RHS have incompatible shapes. Expected RHS to have"
        " shape (k, n) and accumulator to have shape (m, n). Accumulator:"
        f" {acc.shape}. RHS: {b.shape}."
    )

  if (lhs_k * (1 + is_sparse)) != rhs_k:
    raise ValueError(
        f"LHS and RHS have incompatible shapes. LHS: {a.shape}. RHS: {b.shape}.")

  if isinstance(acc, pallas_core.TransformedRef):
    acc_transforms_leaves, acc_transforms_tree = jax.tree.flatten(
        acc.transforms)
    acc = acc.ref
  else:
    acc_transforms_leaves, acc_transforms_tree = [], None

  if isinstance(a, pallas_core.TransformedRef):
    a_transforms_leaves, a_transforms_tree = jax.tree.flatten(a.transforms)

View on GitHub (pinned to 1e1c6a8fc0)