JAX on TPUs and GPUs

JAX was designed to run efficiently on accelerators from day one. Here's what you need to know about running JAX on different hardware.

JAX on GPUs

import jax

# Check available devices
print(jax.devices())
# [CudaDevice(id=0), CudaDevice(id=1), ...]

# JAX automatically uses GPU if available
x = jax.numpy.ones((1000, 1000))
# x is already on GPU — no .to('cuda') needed!

# Multi-GPU: use sharding (see Track 12)
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

devices = jax.devices()  # all GPUs
mesh = Mesh(devices, ('data',))

# Data parallelism: shard batch across GPUs
data_sharding = NamedSharding(mesh, P('data'))
batch = jax.device_put(x, data_sharding)

# As of JAX 0.6.0: requires CUDA 12.8+
# pip install jax[cuda12]

JAX on TPUs

# On Google Cloud TPU VMs:
# pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

# TPU pods: automatic multi-host setup
print(jax.device_count())  # e.g., 8 for a v4-8, 32 for a v4-32

# TPU-specific tips:
# 1. Use bfloat16 — TPUs have native bfloat16 support
x_bf16 = x.astype(jax.numpy.bfloat16)

# 2. Pad batch sizes to multiples of 128 (TPU-friendly)
# 3. Avoid scalar operations — TPUs are designed for large tensor ops
# 4. Use jax.profiler for TPU-specific profiling

Performance Tips

# Common performance pitfalls:

# 1. UNNECESSARY RECOMPILATION
# BAD: different shapes cause recompilation
for batch in variable_size_batches:
    result = jax.jit(fn)(batch)  # recompiles every new shape!

# GOOD: pad to fixed size
max_batch_size = 256
for batch in batches:
    padded = pad_to_size(batch, max_batch_size)
    result = jax.jit(fn)(padded)  # compiled once, reused

# 2. HOST-DEVICE TRANSFER
# BAD: pulling values back to CPU in a loop
for step in range(1000):
    loss = train_step(params, batch)
    print(float(loss))  # blocks! transfers to CPU every step

# GOOD: only transfer periodically
for step in range(1000):
    loss = train_step(params, batch)
    if step % 100 == 0:
        print(float(loss))  # transfer only every 100 steps

# 3. USE jax.block_until_ready() for timing
import time
x = jax.numpy.ones((1000, 1000))
start = time.time()
y = x @ x
y.block_until_ready()  # wait for computation to finish
print(f"Time: {time.time() - start:.4f}s")