This demo explores how to achieve high-performance AI on Google's Tensor Processing Units (TPUs) using the JAX ecosystem, with a specific focus on image recognition workflows. We’ll begin with micro-benchmarks that showcase JAX's unique advantages for TPU-based computation, such as its single-program, multiple-data (SPMD) programming model, which is ideal for the TPU's systolic array architecture. This setup is designed for integration into larger, production-grade environments, such as those running on Kubernetes.
Nathan Beach
Nathan Beach is Director of Product Management for Google Kubernetes Engine (GKE). He leads the product team working to make GKE a great platform on which to run AI workloads. He received his MBA from Harvard Business School and, prior to Google, led his own startup. He is a builder and creator passionate about making products that superbly meet user needs. He enjoys career coaching and mentoring, and he is eager to help others transition into product management and excel in their careers.