{"record":{"id":"54b2dfc7c963a0d9","repo":"jax-ml/jax","slug":"pallas-tpu-requires-a-recent-libtpu-version-at-le","errorCode":null,"errorMessage":"Pallas TPU requires a recent libtpu version (at least 0.0.46). Found version string:\\n{platform_version}","messagePattern":"Pallas TPU requires a recent libtpu version \\(at least 0\\.0\\.46\\)\\. Found version string:\\\\n(.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"critical","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":1094,"sourceCode":"\ndef lower_jaxpr_into_pipelined_module(\n    lowering_context: mlir.LoweringRuleContext,\n    module: ir.Module,\n    grid_mapping: pallas_core.GridMapping,\n    jaxpr: jax_core.Jaxpr,\n    *,\n    name: str,\n    dimension_semantics: Sequence[tpu_core.DimensionSemantics] | None,\n    kernel_type: tpu_core.CoreType,\n    mesh: mesh_lib.Mesh | None = None,\n    dynamic_shape_replacement_enabled: bool = False,\n    fuse_transposed_lhs_in_matmul: bool = False,\n) -> None:\n  backend = lowering_context.module_context.get_backend(optional=True)\n  # NOTE: We should bump this periodically\n  if not is_libtpu_at_least(\"0.0.46\"):\n    platform_version = xla_bridge.get_backend().platform_version\n    raise RuntimeError(\n        \"Pallas TPU requires a recent libtpu version (at least 0.0.46). Found\"\n        f\" version string:\\n{platform_version}\"\n    )\n  debug_info = jaxpr.debug_info\n  _mosaic_lowering_dynamic_shape_env = None\n  if dynamic_shape_replacement_enabled:\n    _mosaic_lowering_dynamic_shape_env = LoweringDynamicShapeEnv()\n\n    def dynamic_shape_replacement_fn(\n        shape: jax_core.Shape,\n    ) -> tuple[Any, ...]:\n      assert _mosaic_lowering_dynamic_shape_env is not None\n      return tuple(\n          _mosaic_lowering_dynamic_shape_env.to_placeholder(dim_expr)\n          if jax_core.is_dim(dim_expr)\n          else dim_expr\n          for dim_expr in shape\n      )","sourceCodeStart":1076,"sourceCodeEnd":1112,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L1076-L1112","documentation":"Pallas's TPU lowering relies on Mosaic/PJRT features that only exist in libtpu >= 0.0.46. At the start of lower_jaxpr_into_pipelined_module the code checks is_libtpu_at_least('0.0.46') against the installed libtpu's platform_version string and raises RuntimeError if the runtime is older.","triggerScenarios":"Running pallas_call on TPU (or any code path that triggers lower_jaxpr_to_pipelined_module) with an outdated jaxlib/libtpu, e.g. a stale nightly wheel, an old TPU VM image, or a CPU/CPU-simulated environment whose libtpu reports an older version string.","commonSituations":"Upgrading JAX but not jaxlib; pinned old jaxlib versions in requirements.txt; TPU VM images shipping old libtpu; mixing pip-installed jax with a system libtpu; environments where platform_version is empty or unparseable so the version check fails.","solutions":["Upgrade jaxlib to a version bundling libtpu >= 0.0.46 (pip install -U jax jaxlib), keeping jax and jaxlib versions in sync","On TPU VMs, update the VM image or reinstall the matching libtpu nightly wheel","Print jax.lib.xla_bridge.get_backend().platform_version to confirm what libtpu the runtime actually loads and fix shadowed installations","If you don't need Pallas, avoid pallas_call on this environment until libtpu is upgraded"],"exampleFix":"# before\npip install jax==0.4.x  # stale jaxlib still installed\n\n# after\npip install -U jax jaxlib  # libtpu >= 0.0.46 bundled","handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic.linalg_util import is_libtpu_at_least  # or jax._src.pallas.mosaic.lowering\nimport jax\nif not is_libtpu_at_least('0.0.46'):\n    raise RuntimeError('Upgrade jax/jaxlib: libtpu too old for Pallas TPU')","typeGuard":null,"tryCatchPattern":"try:\n    kernel = pallas.tpu_kernel(f)\nexcept RuntimeError as e:\n    if 'recent libtpu' in str(e):\n        os.system('pip install -U jax jaxlib')  # then restart the process","preventionTips":["Pin jax and jaxlib to matching recent versions","On TPU VMs, keep the VM image / libtpu nightly current","Log platform_version at startup to detect stale libtpu early"],"tags":["jax","pallas","tpu","libtpu","version-mismatch"],"backgroundTag":"jaxlib-libtpu-version-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}