jax-ml/jax · error · ValueError

The minor dimension size of an accumulator ref must be {info

Error message

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

What it means

The accumulator's last (minor) dimension must exactly equal the MXU column size reported by TPU info (typically the lane width). Any other minor dimension cannot be laid out in accumulator hardware.

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)

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 {shape}
• The product of the major dimensions must be a multiple of {i
• Acc ref dtype must be float32 or int32, got {dtype}
• Accumulators are not available on TPU {info.chip_version}
• mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}

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