jax-ml/jax · error · ValueError

Accumulator and RHS have incompatible shapes. Expected RHS t

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

In non-collective tcgen05.mma (collective_axis=None), the accumulator's N must equal the RHS's n: acc.shape[1] != b.shape[1].

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)

Solutions

1. Match accumulator N to b.shape[1] exactly
2. If porting from collective mode, halve the accumulator N
3. 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)

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

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

• Accumulator and LHS have incompatible shapes. Expected LHS t
• TMEM aliasing only supported for Refs with the same first di
• Accumulator and RHS have incompatible shapes. Expected RHS t
• MMA rhs swizzle must match lhs swizzle. {lhs_swizzle=} {rhs_
• Block shape for {origin} (= {block_shape}) must have the sam

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