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 * 2) in collective mode. Accumulator: {acc.shape}. RHS: {b.shape}.

What it means

When collective_axis is set on tcgen05.mma, the collective MMA produces an N dimension twice the RHS's n, so the accumulator must have shape (m, n*2). This fires when acc.shape[1] != b.shape[1] * 2.

Source

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

    collective_axis: The name of the cluster axis along which to perform
      a collective MMA. The cluster axis should have a size of exactly 2,
      and must be on the minormost cluster axis.
  """
  acc_m, acc_n = acc.shape
  lhs_m, lhs_k = a.shape
  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate/size the accumulator N as 2 * b.shape[1] when using collective_axis
2. Double-check the collective MMA docs for the (m, n*2) accumulator layout
3. If you don't need 2CTA tensor-core mode, drop collective_axis

Example fix

# before
acc = allocate(TMEM, (m, b.shape[1]), jnp.float32)
tcgen05.mma(a, b, acc, k_dim=k, collective_axis=0)
# after
acc = allocate(TMEM, (m, 2 * b.shape[1]), jnp.float32)
tcgen05.mma(a, b, acc, k_dim=k, collective_axis=0)

n_acc = 2 * b.shape[1] if collective_axis is not None else b.shape[1]
assert acc.shape == (a.shape[0], n_acc)
tcgen05.mma(a, b, acc, k_dim=a.shape[1], collective_axis=collective_axis)

When it happens

Trigger: Using tcgen05.mma(..., collective_axis=0) with an accumulator sized (m, n) matching b's n instead of (m, 2*n).

Common situations: Switching a kernel from non-collective to collective MMA without resizing the accumulator; porting examples that omit the doubled N in collective mode.

Related errors

• TMEM aliasing only supported for Refs with the same first di
• Accumulator and LHS have incompatible shapes. Expected LHS t
• Accumulator and RHS have incompatible shapes. Expected RHS t
• Accumulator Ref must be collective if collective_axis is set
• LHS Ref must be collective if collective_axis is set.

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