jax-ml/jax · error · ValueError

Unsupported TPU device kind: {device_kind}. If you are not r

Error message

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` argument specifies the exact TPU version you intend to target.

What it means

get_tpu_info could not resolve the current device kind to TPU info: the kind is neither parseable as a chip version nor in the static registry. On non-TPU hardware you must declare the target TPU via an abstract mesh.

Source

Thrown at jax/_src/tpu_info.py:573

          ),
      )
    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` "
        "argument specifies the exact TPU version you intend to target."
    )
  return _get_tpu_info_impl(chip_version, get_num_device_cores())


@jax_util.cache(trace_context_in_key=True)
def get_tpu_info_for_chip(
    chip_version: ChipVersion, num_tensor_cores_per_logical_device: int
) -> TpuInfo:
  """Returns the TPU hardware info for the given TPU chip version.

  Note that all information is *per-TensorCore* so you would need to multiply by
  `num_tensor_cores_per_logical_device` to obtain the total for the chip.

  Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Wrap code in jax.sharding.use_abstract_mesh with an AbstractMesh specifying the TPU version (e.g. 'TPU v5e')
  2. Run on an actual supported TPU
  3. Upgrade JAX for new device kind support

Example fix

# before
y = jax.lax.top_k(x, k)
# after
import jax.sharding as jsh
with jsh.use_abstract_mesh(jsh.AbstractMesh(1, 'i', axis_types=(jsh.AxisType.Auto,), devices_json=None), names=()):
  # or construct AbstractMesh targeting TPU v5e
  y = jax.lax.top_k(x, k)
Defensive patterns

Strategy: fallback

Validate before calling

import jax
if jax.devices()[0].platform != 'tpu':
    mesh = jax.sharding.AbstractMesh(... , names=..., axis_types=...)  # target e.g. TPU v5e
    cm = jax.sharding.use_abstract_mesh(mesh)

Type guard

def on_tpu() -> bool:
    return jax.devices()[0].platform == 'tpu'

Try / catch

try:
    out = jax.lax.top_k(x, k)
except ValueError as e:
    if 'Unsupported TPU device kind' in str(e):
        out = fallback_top_k(x, k)
    else:
        raise

Prevention

When it happens

Trigger: Calling TPU-specific lowerings (e.g. jax.lax.top_k, logistic, accumulator ops) on CPU/GPU without specifying a target TPU version, or running on an unrecognized TPU kind.

Common situations: Unit-testing TPU-only code paths locally; new/unknown device kinds; CI on CPU.

Understand the failure class

Background: "environment variable is not set" and "Missing keys in environment" errors: what missing required env var messages mean and how to fix them — this error's family across 28 libraries.

Related errors

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