jax-ml/jax · error · ValueError

The shape of the accumulator {acc_shape} does not match the

Error message

The shape of the accumulator {acc_shape} does not match the shape of the lhs {lhs.shape}.

What it means

In matmul_acc_lhs, the leading dimension M of the 2D accumulator must equal lhs.shape[0], because each matmul step accumulates lhs @ rhs into the acc rows. A mismatch is rejected at abstract-eval time.

Source

Thrown at jax/_src/pallas/mosaic/primitives.py:1303

@matmul_acc_lhs_p.def_effectful_abstract_eval
def _matmul_acc_lhs_abstract_eval(
    acc: state.AbstractRef, lhs, *flat_acc_transforms, acc_transforms_tree, load_staged_rhs
):
  del load_staged_rhs,  # Unused.
  transforms = tree_util.tree_unflatten(acc_transforms_tree, flat_acc_transforms)
  if not isinstance(acc.memory_space, tpu_core.AccMemorySpace):
    raise ValueError(f"Expected an accumulator ref, got {acc}")
  transformed_acc = state.transform_type(transforms, acc)
  assert isinstance(transformed_acc, state.AbstractRef)
  acc_shape: tuple[int, ...] = transformed_acc.shape
  if len(acc_shape) != 2:
    raise ValueError(
        f"The shape of the accumulator {acc_shape} is not 2-dimensional."
    )
  m, _ = acc_shape
  if m != lhs.shape[0]:
    raise ValueError(
        f"The shape of the accumulator {acc_shape} does not "
        f"match the shape of the lhs {lhs.shape}."
    )
  return [], {mxu_effect, state.ReadEffect(0), state.WriteEffect(0)}


matmul_pop_p = jax_core.Primitive("matmul_pop")


def matmul_pop(acc: Ref) -> jax.Array:
  """Returns the result of a matrix multiplication from a specific MXU and zeroes the accumulator.

  If the result is not ready yet (the MXU is still busy), the operation blocks.

  ```{warning}
  The kernel must not leave any data in the accumulator upon exit.
  ```

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make lhs block M equal acc.shape[0] (both typically 128 on TPU MXU)
Recompute tiling so lhs blocks are [acc_m, k]

Example fix

# before
matmul_acc_lhs(acc, lhs)  # acc (128,512), lhs (64,512)
# after
lhs = lhs.reshape(2, 64, 512)  # loop over 64-row slabs, or reallocate acc with M=64

Defensive patterns

Strategy: validation

Validate before calling

assert acc_shape[0] == lhs.shape[0], 'M dims must match'

Type guard

def matmul_shapes_compatible(acc_shape, lhs_shape) -> bool:
    return len(acc_shape) == 2 and acc_shape[0] == lhs_shape[0]

Prevention

Derive tiling constants (e.g. 128) from one shared config

When it happens

Trigger: Calling matmul_acc_lhs(acc, lhs) where acc.shape[0] != lhs.shape[0], e.g. acc of shape (128, 512) with lhs of shape (64, 512).

Common situations: Tiling lhs with a block size that differs from the accumulator's M dimension; mixing tile sizes between the lhs pipeline and the acc allocation.

Related errors

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