tpu
**Google TPU (Tensor Processing Unit)**
**What is TPU?**
Purpose-built ASIC for ML training and inference, available via Google Cloud.
**TPU Versions**
| Version | Year | Features |
|---------|------|----------|
| TPU v2 | 2017 | 180 TFLOPS |
| TPU v3 | 2018 | 420 TFLOPS, liquid cooled |
| TPU v4 | 2021 | 275 TFLOPS, 4096-chip pods |
| TPU v5e | 2023 | Cost-optimized inference |
| TPU v5p | 2023 | Training-optimized |
**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**
```python
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**
```python
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**
```python
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**
| Aspect | TPU | GPU (H100) |
|--------|-----|------------|
| Best for | Google ecosystem | General |
| Memory | 16-64GB HBM | 80GB HBM |
| Interconnect | TPU pods | NVLink |
| Software | JAX/TF | PyTorch/TF |
| Availability | GCP only | Universal |
**Pricing (GCP)**
| Type | On-demand | Spot |
|------|-----------|------|
| TPU v4 | $3.22/hr | $0.97/hr |
| TPU v5e | $1.20/hr | $0.36/hr |
**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