jax-ml/jax · error · ValueError

Accumulator and LHS have incompatible shapes. Expected LHS…

Error message

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

What it means

Shape guard in tcgen05_mma: with an accumulator of shape (m, n), the lhs operand must have leading dimension m (overall shape (m, k)). The error fires when acc_m != lhs_m, i.e., the 'a' operand's row count doesn't match the accumulator's, so the MMA cannot tile cleanly into acc.

Solutions

1. Size the accumulator's first dimension to match the LHS block's first dimension
2. Verify the BlockMapping/scratch shapes that allocate the TMEM accumulator
3. Re-derive accumulator shape as (a.shape[0], b.shape[1]) (or n*2 in collective mode)

Example fix

# before
acc = allocate(TMEM, (64, n), jnp.float32)
tcgen05.mma(a, b, acc, k_dim=16)
# after
acc = allocate(TMEM, (a.shape[0], n), jnp.float32)
tcgen05.mma(a, b, acc, k_dim=16)

m, k = a.shape
assert acc.shape[0] == m, f'acc m {acc.shape[0]} != lhs m {m}'
tcgen05.mma(a, b, acc, k_dim=k)

When it happens

Trigger: Calling tcgen05.mma with an accumulator allocated as (m2, n) where m2 differs from the m rows of the LHS block, e.g. LHS (128, 16) with accumulator (64, 64).

Common situations: Mismatched BlockMapping causing the accumulator grid dimension to differ from the LHS block's M; manually allocated TMEM accumulator sized for a different tile shape.

Related errors

• Accumulator and RHS have incompatible shapes. Expected RHS…
• Accumulator and RHS have incompatible shapes. Expected RHS…
• MMA rhs swizzle must match lhs swizzle.
• TMEM aliasing only supported for Refs with the same first…
• a_scale and b_scale must both be present or absent.

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

Appendix: source

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

      Must have orders_tensor_core set to True. If not specified, the MMA
      completion should be explicitly observed by calling
      :func:`jax.experimental.pallas.mosaic_gpu.tcgen05_commit_arrive`
    a_scale: An optional scale for the ``a`` operand. Must be a TMEM Ref if present.
    b_scale: An optional scale for the ``b`` operand. Must be a TMEM Ref if present.
    a_sparse_metadata: An optional sparse metadata for the ``a`` operand.
      Must be a TMEM Ref if present.
    accumulate: Whether to accumulate into acc or overwrite it.
    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}."
    )

View on GitHub (pinned to 1e1c6a8fc0)