Google TPU (Tensor Processing Unit)

What is TPU? Purpose-built ASIC for ML training and inference, available via Google Cloud.

TPU Versions

TPU Architecture

• Matrix multiply units (MXUs) for matmul
• High-bandwidth memory (HBM)
• Interconnect for multi-chip scaling
• Optimized for BF16/INT8

Using TPUs with JAX

import jax
import jax.numpy as jnp

# Check TPU availability
print(jax.devices())  # [TpuDevice(...)]

# Arrays automatically use TPU
x = jnp.ones((1000, 1000))
y = jnp.dot(x, x)

Multi-TPU Training

from jax.sharding import PartitionSpec, NamedSharding
from jax.experimental import mesh_utils

# Create device mesh
devices = mesh_utils.create_device_mesh((4, 2))  # 4x2 TPU grid

# Shard data across devices
mesh = Mesh(devices, axis_names=("data", "model"))
sharding = NamedSharding(mesh, PartitionSpec("data", None))

# Distribute array
distributed_data = jax.device_put(data, sharding)

TPU with TensorFlow

import tensorflow as tf

# TPU initialization
resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

# Create strategy
strategy = tf.distribute.TPUStrategy(resolver)

with strategy.scope():
    model = create_model()
    model.compile(...)
model.fit(dataset)

TPU vs GPU Comparison

Pricing (GCP)

Best Practices

• Use JAX or TensorFlow for best support
• PyTorch works via torch-xla
• Consider v5e for inference (cost-effective)
• Use pods for large model training
• Monitor utilization via Cloud Console