ABOUT EVENT

Workshop Description

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.

Prerequisites

Access to KOA, comfortability with UNIX systems, math proficiency, machine learning basics, python 3 installed and proficiency. https://github.com/jax-ml/jax

Learning Objectives

This workshop aims to provide attendees with an overview of the following: Mathematics of implicit differentiation and why it is used to compute gradients through optimization solvers.

Implement optimization problems as differentiable layers within end-to-end machine learning and scientific pipelines.

Wield JAX’s core transformations. jit for compilation, grad for automatic differentiation, and vmap for vectorization.

Configure, profile, and scale JAX workloads across the distributed GPU nodes of the Koa HPC cluster.”

Tools Used

• Python
• NumPy
• Slurm
• JAX

Registration: https://ci.its.hawaii.edu/portal/workshops/hidsi-workshop-series/differentiable-optimization/