ABOUT EVENT

This hands-on workshop covers how to use JAX on the Koa HPC cluster to implement differentiable optimization. This approach bridges scientific computing and machine learning by letting you insert entire optimization solvers as layers inside end-to-end computational models. Instead of unrolling iterative solvers in memory, we will use the Implicit Function Theorem to compute exact gradients directly from the final optimal solution. This mathematical shortcut lets you backpropagate through complex physical constraints, simulations, or engineering objectives with minimal memory overhead.

We will focus on implementing and scaling these workflows using JAX’s core features: automatic differentiation (grad), vectorization (vmap), and XLA compilation (jit). You will learn how to write custom optimization layers, profile your code for bottlenecks, and deploy distributed jobs across Koa’s GPU nodes.