Accelerating AI TPU & JAX: Hands-On Workshop

Join us for an intensive hands-on technical session hosted by TFUG Prayagraj, where we dive into Google’s high-performance hardware and software stack: Tensor Processing Units (TPUs) and JAX. Participants will learn how custom AI hardware works and gain practical experience running high-speed numerical operations, model training, and automatic differentiation using Google Colab's free Cloud TPU runtimes.

#### Detailed Agenda

| Time | Session / Activity | Description |
| ---- | ------------------ | ----------- |
| 10:00 AM – 10:30 AM | Check-In & Environment Setup | Arrival, attendee check-in, setting up Google Colab runtime, and enabling Cloud TPU v2-8 / v3-8 backends. |
| 10:30 AM – 11:15 AM | Keynote: Why TPU + JAX? | Overview of Google AI hardware evolution (GPU vs. TPU architecture), XLA compiler, and introduction to the JAX ecosystem. |
| 11:15 AM – 12:15 PM | Session 1: JAX Fundamentals on Colab | Hands-on: NumPy vs. JAX (`jax.numpy`), pure functions, static tracking, and automatic differentiation using `jax.grad`. |
| 12:15 PM – 01:15 PM | Session 2: Unlocking TPU Speed | Hands-on: Compiling code with `jax.jit`, vectorizing functions with `jax.vmap`, and benchmarking CPU vs. GPU vs. TPU execution times. |
| 01:15 PM – 02:00 PM | Lunch & Networking Break | Networking, peer discussions, and Q&A with GDG mentors. |
| 02:00 PM – 03:15 PM | Session 3: Multi-Core TPU Parallelization | Hands-on Lab: Distributing tensor computations across all 8 TPU cores using `jax.pmap` / `jax.experimental.shard_map` on a neural network training task. |
| 03:15 PM – 03:45 PM | Post-Event Roadmap & Guided Problem Set | Introduction to offline practice worksheets, self-paced Kaggle/Colab problem statements, and extended learning resources. |
| 03:45 PM – 04:00 PM | Closing, DEV Badges & Group Photo | Evaluation of lab completion, awarding digital badges, and official photo session. |

***

# Speakers

### Atul

Session: Google Colab

###

###