jax-ml/jax · error · ValueError

Lite chips, single core chips, and dual-core chips that do n

Error message

Lite chips, single core chips, and dual-core chips that do not support Megacore must have num_tensor_cores_per_logical_device=1, but got {num_tensor_cores_per_logical_device}.

What it means

get_tpu_info_for_chip validates that lite, single-core, and non-Megacore dual-core chips must have num_tensor_cores_per_logical_device=1, but a different value was supplied.

Source

Thrown at jax/_src/tpu_info.py:612

      device in the requested configuration. Should be 1 for single-core chips
      (TPU_V4I, TPU_V5E, TPU_V6E). For dual-core chips that support Megacore
      (TPU_V4, TPU_V5P), this can be 2 (Megacore mode) or 1 (split mode). For
      dual-core chips that do not support Megacore (TPU_V2, TPU_V3, TPU_7X),
      this must be 1.
  """
  if (
      chip_version.is_lite
      or chip_version
      in {
          ChipVersion.TPU_V2,
          ChipVersion.TPU_V3,
          ChipVersion.TPU_7,
          ChipVersion.TPU_7X,
          ChipVersion.TPU_8I,
          ChipVersion.TPU_8T,
      }
  ) and num_tensor_cores_per_logical_device != 1:
    raise ValueError(
        "Lite chips, single core chips, and dual-core chips that do not support"
        " Megacore must have num_tensor_cores_per_logical_device=1, but got"
        f" {num_tensor_cores_per_logical_device}."
    )

  return _get_tpu_info_impl(chip_version, num_tensor_cores_per_logical_device)


# TODO(sharadmv): Generalize Tiling to capture the various options
# (compact 2nd minor, large 2nd minor, regular tiling)
class Tiling(enum.Enum):
  COMPACT = enum.auto()
  SPARSE_CORE = enum.auto()

  @property
  def shape(self) -> tuple[int, ...]:
    # TODO(slebedev): Use ``get_tpu_info()`` instead of hardcoding the values.
    match self:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass num_tensor_cores_per_logical_device=1 for such chips
2. Use the default device-derived value instead of a manual one

Example fix

# before
info = get_tpu_info_for_chip(ChipVersion.TPU_5E, 2)
# after
info = get_tpu_info_for_chip(ChipVersion.TPU_5E, 1)

SINGLE_CORE = {ChipVersion.TPU_4_LITE, ChipVersion.TPU_5E, ...}
if chip in SINGLE_CORE:
    num_tensor_cores_per_logical_device = 1

When it happens

Trigger: Calling get_tpu_info_for_chip(chip, num_tensor_cores_per_logical_device=2) for chips like TPU v4 lite / v5e / single-core variants.

Common situations: Manually constructing TpuInfo for testing or simulation with wrong core counts; copy-pasting v4 (Megacore) config.

Related errors

• Invalid memory space: {memory_space!r}
• number of KV heads must be even when megacore_mode is 'kv_he
• batch size must be even when megacore_mode is 'batch'
• megacore_mode must be one of ['kv_head', 'batch', None]
• {num_queries_per_block=} must be positive.

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