jax-ml/jax · error · ValueError

mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}

Error message

mxu_id must be in [0, {info.num_mxus}), got {mxu_id=}

What it means

Accumulator refs are per-MXU resources; mxu_id must index a valid MXU on the chip (0 <= mxu_id < info.num_mxus). Values outside that range (including negatives) fail validation.

Source

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

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}"
    )


class MemoryRef(pallas_core.MemoryRef):

  def __matmul__(self, other, /):
    if not isinstance(other, pallas_core.Mesh):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Query info.num_mxus via tpu_info.get_tpu_info() and clamp/validate mxu_id against it
2. Use device-derived indices rather than hardcoded constants
3. Fix off-by-one loops generating mxu_id values

Example fix

// before
acc = make_acc(..., mxu_id=8)
// after
info = tpu_info.get_tpu_info()
assert 0 <= mxu_id < info.num_mxus
acc = make_acc(..., mxu_id=mxu_id % info.num_mxus)

from jax._src.pallas.mosaic import tpu_info
info = tpu_info.get_tpu_info()
assert 0 <= mxu_id < info.num_mxus, (mxu_id, info.num_mxus)

When it happens

Trigger: Passing an out-of-range mxu_id when creating an accumulator ref, e.g., mxu_id=4 on a chip with 2 MXUs, or defaulting to -1 as 'unspecified'.

Common situations: Hardcoding mxu_id values tuned for one TPU generation and running on another; looping over core indices with off-by-one errors.

Related errors

• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}
• Accumulators are not available on TPU {info.chip_version}
• 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/3f258f2ba07fb938. Report an issue: GitHub.