{"record":{"id":"629aaca440859592","repo":"jax-ml/jax","slug":"unsupported-tpu-device-kind-device-kind-if-you","errorCode":null,"errorMessage":"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.","messagePattern":"Unsupported TPU 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\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu_info.py","lineNumber":573,"sourceCode":"          ),\n      )\n    case _:\n      raise ValueError(f\"Unsupported TPU chip version: {chip_version}\")\n\n\n@jax_util.cache(trace_context_in_key=True)\ndef get_tpu_info() -> TpuInfo:\n  \"\"\"Returns the TPU hardware info for the current device.\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  device_kind = get_device_kind()\n  chip_version = chip_version_from_device_kind(device_kind)\n  if chip_version is None:\n    if device_kind in registry:\n      return registry[device_kind]()\n    raise ValueError(\n        f\"Unsupported TPU device kind: {device_kind}. If you are not running \"\n        \"on a TPU device, you need to wrap your code in a \"\n        \"`jax.sharding.use_abstract_mesh` context manager whose `AbstractMesh` \"\n        \"argument specifies the exact TPU version you intend to target.\"\n    )\n  return _get_tpu_info_impl(chip_version, get_num_device_cores())\n\n\n@jax_util.cache(trace_context_in_key=True)\ndef get_tpu_info_for_chip(\n    chip_version: ChipVersion, num_tensor_cores_per_logical_device: int\n) -> TpuInfo:\n  \"\"\"Returns the TPU hardware info for the given TPU chip version.\n\n  Note that all information is *per-TensorCore* so you would need to multiply by\n  `num_tensor_cores_per_logical_device` to obtain the total for the chip.\n\n  Args:","sourceCodeStart":555,"sourceCodeEnd":591,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu_info.py#L555-L591","documentation":"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.","triggerScenarios":"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.","commonSituations":"Unit-testing TPU-only code paths locally; new/unknown device kinds; CI on CPU.","solutions":["Wrap code in jax.sharding.use_abstract_mesh with an AbstractMesh specifying the TPU version (e.g. 'TPU v5e')","Run on an actual supported TPU","Upgrade JAX for new device kind support"],"exampleFix":"# before\ny = jax.lax.top_k(x, k)\n# after\nimport jax.sharding as jsh\nwith jsh.use_abstract_mesh(jsh.AbstractMesh(1, 'i', axis_types=(jsh.AxisType.Auto,), devices_json=None), names=()):\n  # or construct AbstractMesh targeting TPU v5e\n  y = jax.lax.top_k(x, k)","handlingStrategy":"fallback","validationCode":"import jax\nif jax.devices()[0].platform != 'tpu':\n    mesh = jax.sharding.AbstractMesh(... , names=..., axis_types=...)  # target e.g. TPU v5e\n    cm = jax.sharding.use_abstract_mesh(mesh)","typeGuard":"def on_tpu() -> bool:\n    return jax.devices()[0].platform == 'tpu'","tryCatchPattern":"try:\n    out = jax.lax.top_k(x, k)\nexcept ValueError as e:\n    if 'Unsupported TPU device kind' in str(e):\n        out = fallback_top_k(x, k)\n    else:\n        raise","preventionTips":["Wrap CPU unit tests of TPU ops in use_abstract_mesh","Provide non-TPU fallbacks for TPU-only primitives"],"tags":["tpu","cpu-testing","abstract-mesh","device-kind","jax"],"backgroundTag":"missing-env-var","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}