jax-ml/jax · error · ValueError

The minor dimension size of an accumulator ref must be

Error message

The minor dimension size of an accumulator ref must be {info.mxu_column_size} but got {n}

What it means

Validation in check_accumulator_ref: an accumulator ref's minor (last) dimension must equal the TPU MXU column size (info.mxu_column_size), since accumulators live in MXU columns; a ref whose last dimension n differs cannot be mapped onto the hardware accumulator and raises ValueError.

Solutions

1. Set shape[-1] to exactly info.mxu_column_size (query it at runtime)
2. Pad the N dimension of your GEMM tile up to the MXU column size and slice after loading the accumulator
3. Regenerate tile shapes per target TPU generation instead of hardcoding

Example fix

// before
acc = alloc_acc((M, 100))
// after
info = tpu_info.get_tpu_info()
acc = alloc_acc((M, info.mxu_column_size))
result = acc.load()[:, :100]

from jax._src.pallas.mosaic import tpu_info
info = tpu_info.get_tpu_info()
assert shape[-1] == info.mxu_column_size, (shape, info.mxu_column_size)

When it happens

Trigger: Allocating an ACC ref whose shape[-1] differs from info.mxu_column_size, e.g. shape (128, 100) on hardware with 128-wide MXU columns requiring exactly 128.

Common situations: Using non-power-of-two or arbitrary N dimensions in matmul tiles; porting kernels between TPU generations with different MXU widths.

Related errors

• Acc ref must be at least 2D, got shape
• The product of the major dimensions must be a multiple of
• Acc ref dtype must be float32 or int32, got
• Accumulators are not available on TPU
• Expected -dim 'key' tensor for MQA. Instead got a -dim one.

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

Appendix: source

Thrown at jax/_src/pallas/mosaic/core.py:232

  if len(shape) < 2:
    raise ValueError(f"Acc ref must be at least 2D, got shape {shape}")

  if dtype not in (jnp.float32, jnp.int32):
    raise ValueError(
        f"Acc ref dtype must be float32 or int32, got {dtype}")

  info = tpu_info.get_tpu_info()
  if not info.num_accumulators:
    raise ValueError(
        f"Accumulators are not available on TPU {info.chip_version}"
    )

  if mxu_id < 0 or mxu_id >= info.num_mxus:
    raise ValueError(f"mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}")

  m, n = math.prod(shape[:-1]), shape[-1]
  if n != info.mxu_column_size:
    raise ValueError(
        f"The minor dimension size of an accumulator ref must be "
        f"{info.mxu_column_size} but got {n}"
    )
  if m <= 0 or m % info.num_sublanes != 0:
    raise ValueError(
        f"The product of the major dimensions must be a multiple of "
        f"{info.num_sublanes}, but got {m}"
    )


class MemoryRef(pallas_core.MemoryRef):

  def __matmul__(self, other, /):
    if not isinstance(other, pallas_core.Mesh):
      return NotImplemented
    return dataclasses.replace(self, memory_space=self.memory_space @ other)

View on GitHub (pinned to 1e1c6a8fc0)