{"record":{"id":"a28ed4f3daf94a3b","repo":"jax-ml/jax","slug":"tpu-generation-is-not-supported","errorCode":null,"errorMessage":"TPU generation is not supported","messagePattern":"TPU generation is not supported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu_info.py","lineNumber":294,"sourceCode":"  def get_sublane_tiling(self, dtype: dtypes.DType) -> int:\n    \"\"\"Returns the sublane tiling for the given itemsize.\n\n    Note that this is a heurustic and depends on the settings of the XLA flags.\n    \"\"\"\n    bitwidth = dtypes.itemsize_bits(dtype)\n    if self.generation < 7:\n      # Caveat: before TPU7x, by default XLA does not use large 2nd minor tiling\n      # but it can be enabled by setting the flag\n      # xla_tpu_enable_large_2nd_minor_layout_for_x16.\n      if bitwidth == 16 or bitwidth == 32:\n        return self.num_sublanes\n      else:\n        # Large 2nd minor tiling is enabled for other types.\n        return self.num_sublanes * (32 // bitwidth)\n    # XLA allows large 2nd minor tiling by default starting with TPU7x.\n    if self.generation == 7 or self.generation == 8:\n      return self.num_sublanes * (32 // bitwidth)\n    raise NotImplementedError(\"TPU generation is not supported\")\n\n\ndef is_tpu_device() -> bool:\n  return chip_version_from_device_kind(get_device_kind()) is not None\n\n\nregistry: dict[str, Callable[[], TpuInfo]] = {}\n\n\ndef _get_tpu_info_impl(chip_version: ChipVersion, num_cores: int) -> TpuInfo:\n  \"\"\"Returns the TPU hardware info for the given chip version and core count.\n\n  Note that all information is *per-TensorCore* so you would need to multiply by\n  `num_cores` to obtain the total for the chip.\n\n  Args:\n    chip_version: The TPU chip version.\n    num_cores: The number of TensorCores per chip for this configuration. This","sourceCodeStart":276,"sourceCodeEnd":312,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu_info.py#L276-L312","documentation":"TpuInfo.get_sublane_tiling was called for a TPU generation (other than 5/5e/6/7/8 handled above) whose sublane tiling rules are unknown to this JAX version.","triggerScenarios":"Running on a newer TPU generation (e.g. TPU9+) with an older JAX that lacks tiling rules, or a corrupted generation parse.","commonSituations":"Cloud TPU images with newer chips than the installed JAX supports.","solutions":["Upgrade JAX/jaxlib to a version supporting your TPU generation","Pin workloads to supported generations"],"exampleFix":null,"handlingStrategy":"fallback","validationCode":"from jax._src.tpu_info import TpuInfo\nassert info.generation in (5,6,7,8)","typeGuard":null,"tryCatchPattern":"try:\n    tiling = info.get_sublane_tiling(bw)\nexcept NotImplementedError:\n    tiling = None  # compute fallback","preventionTips":["Upgrade JAX before running on new TPU generations"],"tags":["tpu","hardware-generation","version-mismatch","jax"],"backgroundTag":"unsupported-hardware-version","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}