jax-ml/jax · error · ValueError

Acc ref must be at least 2D, got shape {shape}

Error message

Acc ref must be at least 2D, got shape {shape}

What it means

Mosaic TPU accumulator refs (pltpu.ACC) map directly onto TPU MXU accumulator hardware, which is inherently 2D (sublanes x lanes). A ref whose shape has fewer than 2 dimensions cannot be backed by that hardware, so check_accumulator_ref raises ValueError immediately.

Source

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

        self, "shape_invariant_numerics", shape_invariant_numerics
    )
    object.__setattr__(self, "use_tc_tiling_on_sc", use_tc_tiling_on_sc)
    object.__setattr__(self, "needs_layout_passes", needs_layout_passes)
    object.__setattr__(
        self,
        "fuse_transposed_lhs_in_matmul",
        fuse_transposed_lhs_in_matmul,
    )
    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 "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Give the accumulator at least 2 dimensions, e.g. reshape (K,) to (1, K) or use the intended (M, N) tile shape
2. Verify the minor dimension equals the MXU column size (see related checks) after fixing rank
3. If you only need a 1D result, compute into a 2D acc and reshape/squeeze after accumulator_load

Example fix

// before
acc = pltpu.ACC  used with shape (256,)
// after
acc_shape = (1, 256)  # 2D accumulator; reshape after load
acc = ... allocate with acc_shape

def valid_acc_shape(shape):
    return len(shape) >= 2

def is_2d_plus(shape: tuple[int, ...]) -> TypeGuard[tuple[int, int, ...]]: return len(shape) >= 2

When it happens

Trigger: Creating a Mosaic TPU kernel accumulator ref (e.g., via TensorCore compute tile plumbing or pltpu accumulator APIs) with a 1D or scalar shape, verified in __post_init__/__call__ of the wrapper class.

Common situations: Writing a TPU matmul kernel and allocating the accumulator with shape (N,) or () instead of (M, N); generalizing a GPU Pallas kernel whose accumulator was 1D.

Related errors

• The minor dimension size of an accumulator ref must be {info
• 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/a934a18b3d18a4a5. Report an issue: GitHub.