jax-ml/jax · error · ValueError

Accumulators are not available on TPU {info.chip_version}

Error message

Accumulators are not available on TPU {info.chip_version}

What it means

The target TPU chip reports zero available accumulators (tpu_info.num_accumulators == 0), meaning the hardware generation does not expose MXU accumulators to Mosaic kernels. Accumulator refs therefore cannot be validated further and are rejected.

Source

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

    )
    object.__setattr__(self, "opt_level", opt_level)

  # Replace is a method, not a field.
  replace = dataclasses.replace


def check_accumulator_ref(shape: tuple[int, ...], dtype: jnp.dtype, mxu_id: int):
  from jax._src.pallas.mosaic import tpu_info  # pyrefly: ignore[missing-module-attribute]
  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}"
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check tpu_info.get_tpu_info() output for your device and pick a kernel variant without accumulators
2. Use plain VMEM buffers and regular dot ops instead of ACC refs on this chip
3. Run on a TPU version that exposes accumulators if the kernel requires them

from jax._src.pallas.mosaic import tpu_info
info = tpu_info.get_tpu_info()
if not info.num_accumulators:
    use_vmem_fallback = True

When it happens

Trigger: Running a kernel using ACC refs on a TPU chip version whose info reports no accumulators (older/edge TPU generations or certain emulated environments).

Common situations: Developing on a different TPU generation than production (e.g., v4 vs v5e/v6e class differences); running on simulators/libtpu runtimes that underreport hardware info.

Related errors

• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}
• mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}
• The minor dimension size of an accumulator ref must be {info
• The product of the major dimensions must be a multiple of {i

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