jax-ml/jax · error · ValueError

Accumulator and LHS have incompatible shapes. Expected LHS t

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

In tcgen05.mma, the accumulator's leading dimension m must equal the LHS's leading dimension m (LHS is (m, k), accumulator is (m, n)). This check fires when acc.shape[0] != a.shape[0].

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)

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 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/074017cf68dff23e. Report an issue: GitHub.