jax-ml/jax · error · ValueError

The product of the major dimensions must be a multiple of {i

Error message

The product of the major dimensions must be a multiple of {info.num_sublanes}, but got {m}

What it means

The product of the accumulator's major dimensions (all dims except the last) must be a positive multiple of num_sublanes, because accumulator rows are distributed across TPU sublanes. Zero or non-multiple totals are rejected.

Source

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

        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)


class MemorySpace(enum.Enum):
  VMEM = "vmem"
  VMEM_SHARED = "vmem_shared"
  SMEM = "smem"
  CMEM = "cmem"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Round the product of major dims up to the next multiple of num_sublanes (pad M)
2. Ensure M > 0 (e.g., a leading dimension of 0 from an empty batch)
3. Query info.num_sublanes and derive tile sizes from it per target chip

Example fix

// before
acc = alloc_acc((100, N))
// after
info = tpu_info.get_tpu_info()
M = ceil(100 / info.num_sublanes) * info.num_sublanes
acc = alloc_acc((M, N))

import math
from jax._src.pallas.mosaic import tpu_info
info = tpu_info.get_tpu_info()
m = math.prod(shape[:-1])
assert m > 0 and m % info.num_sublanes == 0, (m, info.num_sublanes)

When it happens

Trigger: Allocating an ACC ref where prod(shape[:-1]) is 0 or not divisible by info.num_sublanes, e.g., major dim 100 when num_sublanes is 128.

Common situations: Choosing M tile sizes that don't match sublane counts; ragged batch dims flattened into the accumulator.

Related errors

• Acc ref must be at least 2D, got shape {shape}
• The minor dimension size of an accumulator ref must be {info
• 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/09172f41b94f161c. Report an issue: GitHub.