jax-ml/jax · critical · RuntimeError

Pallas TPU requires a recent libtpu version (at least 0.0.46

Error message

Pallas TPU requires a recent libtpu version (at least 0.0.46). Found version string:\n{platform_version}

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:1094

def lower_jaxpr_into_pipelined_module(
    lowering_context: mlir.LoweringRuleContext,
    module: ir.Module,
    grid_mapping: pallas_core.GridMapping,
    jaxpr: jax_core.Jaxpr,
    *,
    name: str,
    dimension_semantics: Sequence[tpu_core.DimensionSemantics] | None,
    kernel_type: tpu_core.CoreType,
    mesh: mesh_lib.Mesh | None = None,
    dynamic_shape_replacement_enabled: bool = False,
    fuse_transposed_lhs_in_matmul: bool = False,
) -> None:
  backend = lowering_context.module_context.get_backend(optional=True)
  # NOTE: We should bump this periodically
  if not is_libtpu_at_least("0.0.46"):
    platform_version = xla_bridge.get_backend().platform_version
    raise RuntimeError(
        "Pallas TPU requires a recent libtpu version (at least 0.0.46). Found"
        f" version string:\n{platform_version}"
    )
  debug_info = jaxpr.debug_info
  _mosaic_lowering_dynamic_shape_env = None
  if dynamic_shape_replacement_enabled:
    _mosaic_lowering_dynamic_shape_env = LoweringDynamicShapeEnv()

    def dynamic_shape_replacement_fn(
        shape: jax_core.Shape,
    ) -> tuple[Any, ...]:
      assert _mosaic_lowering_dynamic_shape_env is not None
      return tuple(
          _mosaic_lowering_dynamic_shape_env.to_placeholder(dim_expr)
          if jax_core.is_dim(dim_expr)
          else dim_expr
          for dim_expr in shape
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade jaxlib to a version bundling libtpu >= 0.0.46 (pip install -U jax jaxlib), keeping jax and jaxlib versions in sync
2. On TPU VMs, update the VM image or reinstall the matching libtpu nightly wheel
3. Print jax.lib.xla_bridge.get_backend().platform_version to confirm what libtpu the runtime actually loads and fix shadowed installations
4. If you don't need Pallas, avoid pallas_call on this environment until libtpu is upgraded

Example fix

# before
pip install jax==0.4.x  # stale jaxlib still installed

# after
pip install -U jax jaxlib  # libtpu >= 0.0.46 bundled

from jax._src.pallas.mosaic.linalg_util import is_libtpu_at_least  # or jax._src.pallas.mosaic.lowering
import jax
if not is_libtpu_at_least('0.0.46'):
    raise RuntimeError('Upgrade jax/jaxlib: libtpu too old for Pallas TPU')

try:
    kernel = pallas.tpu_kernel(f)
except RuntimeError as e:
    if 'recent libtpu' in str(e):
        os.system('pip install -U jax jaxlib')  # then restart the process

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Requires libtpu >= 0.1.0
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}

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