TPU Runner
A Google Cloud TPU job orchestrator that races capacity across compatible regions and TPU types, assigns queued jobs, and automatically retries Spot-preempted runs.
Projects
LE-WM JAX
A JAX port of LeWorldModel with PyTorch and Safetensors checkpoint support, tested for parity on CPU and TPU.
Plasticity
A JAX/TPU pipeline for studying whether later pretraining checkpoints lose plasticity during midtraining and supervised fine-tuning.