Google Cloud TPU Runtime Library

0.0.39 · active · verified Sat Apr 11

The `libtpu` library is a low-level runtime component that provides the interface for Python-based machine learning frameworks (such as JAX and PyTorch/XLA) to communicate with Google Cloud TPUs. It is primarily a dependency managed by these high-level frameworks rather than a library intended for direct user application development. The current version is 0.0.39 and it requires Python >= 3.11. It has no strict release cadence, with updates typically coinciding with changes in underlying TPU infrastructure or integrations with ML frameworks.

Warnings

gotcha `libtpu` is a low-level runtime library primarily used by ML frameworks (JAX, PyTorch/XLA) to interface with Google Cloud TPUs. It is not designed for direct import and high-level application development by end-users. Most users will interact with `libtpu` indirectly through these frameworks.
Fix: Focus on setting up your ML framework (e.g., JAX) correctly for TPU usage; avoid trying to directly manipulate `libtpu`'s internal APIs unless you have a highly specialized use case.
gotcha This library requires a Google Cloud TPU environment to function. It cannot be used for local development without significant, specialized setup (e.g., a software emulator), which is not officially supported for general use.
Fix: Deploy your application to a Google Cloud TPU VM or environment to utilize `libtpu`'s functionality.
gotcha Version compatibility between `libtpu`, JAX/PyTorch/XLA, and the underlying Google Cloud TPU software stack is crucial. Mismatched versions can lead to runtime errors, device not found errors, or unexpected behavior.
Fix: Always follow the official installation instructions for JAX or PyTorch/XLA on TPU, which typically bundle or specify compatible `libtpu` versions. Avoid arbitrary `pip install libtpu` if you are using a framework.
gotcha The package currently requires Python 3.11 or higher. Using older Python versions will result in installation failures or runtime errors.
Fix: Ensure your environment is running Python 3.11 or a later compatible version.
gotcha For some environments or older setups, `libtpu` might rely on specific environment variables (e.g., `LD_LIBRARY_PATH`, `XRT_TPU_CONFIG`) being correctly set to locate its underlying C++ components. While modern framework installations often handle this, it can be a source of 'TPU not found' errors.
Fix: Consult Google Cloud's official documentation for JAX/PyTorch on TPU to ensure all necessary environment configurations are in place. When troubleshooting, verify these variables if other solutions fail.

Install

pip install libtpu Install `libtpu` directly
pip install jax[tpu] Recommended: Install via JAX for TPU support (includes libtpu)

Imports

libtpu
```
import libtpu
```
While `libtpu` can be imported, it is a low-level runtime library. Direct import for accessing user-facing classes or functions is uncommon. Most users will rely on frameworks like JAX or PyTorch/XLA to manage and utilize `libtpu` internally.

Quickstart

This quickstart demonstrates how to verify the presence and configuration of TPU devices using JAX, which is the most common way users interact with the capabilities provided by `libtpu`. Running this code on a Google Cloud TPU instance with `jax[tpu]` installed will confirm `libtpu`'s operational status.

import os

# This quickstart demonstrates checking for TPU devices using JAX,
# which implicitly relies on `libtpu` being correctly installed and configured.
# Direct usage of `libtpu`'s API for application logic is rare.

try:
    import jax
    # The following line will only succeed if libtpu is correctly installed
    # and a TPU device is available and configured in the environment.
    tpu_devices = jax.devices('tpu')
    print(f"Found {len(tpu_devices)} TPU devices: {tpu_devices}")
    if tpu_devices:
        print("libtpu is likely working correctly with JAX on a TPU.")
    else:
        print("No TPU devices found. libtpu may be installed, but no TPU is available or configured.")
except ImportError:
    print("JAX is not installed. Please install JAX with TPU support: pip install jax[tpu]")
except Exception as e:
    print(f"An error occurred while checking for TPU devices: {e}")
    print("This could indicate an issue with libtpu installation or TPU environment configuration.")

view raw JSON →