jax-ml/jax · error · ValueError

Acc ref dtype must be float32 or int32, got {dtype}

Error message

Acc ref dtype must be float32 or int32, got {dtype}

What it means

TPU MXU accumulators only support float32 and int32 element types. Passing an accumulator ref with any other dtype (bfloat16, float16, float64, int8, ...) fails this validation before any compilation happens.

Source

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

    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 "
        f"{info.mxu_column_size} but got {n}"
    )
  if m <= 0 or m % info.num_sublanes != 0:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Declare the accumulator as jnp.float32 (standard for matmul accumulation) or jnp.int32
2. Cast to bfloat16 only after loading the accumulator out to VMEM
3. Double-check any dtype inference (e.g., from jnp.zeros default) that might produce non-float32

Example fix

// before
acc = kernel_init(dtype=jnp.bfloat16)
// after
acc = kernel_init(dtype=jnp.float32)
out = acc.load().astype(jnp.bfloat16)

import jax.numpy as jnp
assert acc_dtype in (jnp.float32, jnp.int32), acc_dtype

When it happens

Trigger: Allocating a Mosaic TPU accumulator with jnp.bfloat16 or another dtype instead of jnp.float32/jnp.int32 in check_accumulator_ref (called during wrapper __call__/__post_init__).

Common situations: Writing bf16 matmul kernels (common on TPU) and assuming the accumulator can also be bf16; converting GPU Pallas kernels that accumulate in fp16.

Related errors

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