jax-ml/jax · error · ValueError

The shape of the accumulator {acc_shape} is not 2-dimensiona

Error message

The shape of the accumulator {acc_shape} is not 2-dimensional.

What it means

The MXU accumulator used by matmul_acc_lhs must be a 2D [M, N] tile because the TPU MXU performs 2D matrix products. A 1D or 3D+ accumulator shape is rejected.

Source

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

      *flat_acc_transforms,
      acc_transforms_tree=acc_transforms_treedef,
      load_staged_rhs=load_staged_rhs,
  )


@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.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape/allocate the accumulator as a 2D [M, N] tile matching the MXU dimensions (e.g. 128x128)
2. Handle batching with a loop over batch index, keeping each accumulator 2D

Example fix

# before
acc = new_ref((batch, 128, 128), ..., memory_space=ACC)
# after
for b in range(batch):
  acc = new_ref((128, 128), ..., memory_space=ACC)
  matmul_acc_lhs(acc, lhs[b])

assert len(acc_shape) == 2, f'acc must be 2D, got {acc_shape}'

def is_2d_shape(shape) -> bool:
    return len(tuple(shape)) == 2

When it happens

Trigger: Calling matmul_acc_lhs with an accumulator ref whose shape has rank != 2, e.g. shape (128,) or (8, 128, 128).

Common situations: Allocating the accumulator with the batch dimension left in; using a flat scratch buffer instead of an [M, N] tile.

Related errors

• The shape of the accumulator {acc_shape} does not match the
• Loading from an accumulator is not supported. Use `matmul_po
• Can only load scalars from SMEM
• Expected value and mask to have the same shape, but got valu
• Can only store scalars to SMEM

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