jax-ml/jax · error · ValueError

Unsupported TPU chip version: {chip_version}

Error message

Unsupported TPU chip version: {chip_version}

What it means

chip_version_from_device_kind produced a chip version that _get_tpu_info_impl's match statement has no entry for — a TPU generation unknown to this JAX.

Source

Thrown at jax/_src/tpu_info.py:558

          vmem_capacity_bytes=128 * 1024 * 1024,  # 128 MiB per core
          cmem_capacity_bytes=0,
          smem_capacity_bytes=1024 * 1024,  # 1 MiB per core
          hbm_capacity_bytes=231_000_000_000 // tensor_cores_per_chip,
          mem_bw_bytes_per_second=int(6.4e12 // tensor_cores_per_chip),
          bf16_ops_per_second=int(0.9961e15 // tensor_cores_per_chip),
          int8_ops_per_second=int(0.9961e15 // tensor_cores_per_chip),
          fp8_ops_per_second=int(5.9769e15 // tensor_cores_per_chip),
          int4_ops_per_second=int(11.9538e15 // tensor_cores_per_chip),
          sparse_core=SparseCoreInfo(
              num_cores=2,
              num_subcores=16,
              num_lanes=16,
              vmem_capacity_bytes=256 * 1024,  # 256 KiB per vector subcore
              dma_granule_size_bytes=64,
          ),
      )
    case _:
      raise ValueError(f"Unsupported TPU chip version: {chip_version}")


@jax_util.cache(trace_context_in_key=True)
def get_tpu_info() -> TpuInfo:
  """Returns the TPU hardware info for the current device.

  Note that all information is *per-TensorCore* so you would need to multiply by
  `num_cores` to obtain the total for the chip.
  """
  device_kind = get_device_kind()
  chip_version = chip_version_from_device_kind(device_kind)
  if chip_version is None:
    if device_kind in registry:
      return registry[device_kind]()
    raise ValueError(
        f"Unsupported TPU device kind: {device_kind}. If you are not running "
        "on a TPU device, you need to wrap your code in a "
        "`jax.sharding.use_abstract_mesh` context manager whose `AbstractMesh` "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade JAX/jaxlib to a release supporting the chip
2. File an issue / add the chip entry if you control the fork

from jax._src.tpu_info import ChipVersion
assert chip_version in set(ChipVersion)

When it happens

Trigger: Running on a newly released TPU chip (or a device kind string parsing to an unmapped version) while using an older JAX.

Common situations: New TPU generation rollouts; custom device_kind strings.

Related errors

• __dlpack__ device only supported for TPU pinned host memory
• `buffer_callback` not supported on {platform} backend.
• Partitioned callback not implemented on {platform} backend.
• Custom Partitioning rules must return Sharding.
• Shardy is used, but sharding propagation callbacks instead o

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