PyHEP topical meeting - JAX ecosystem

by Johanna Haffner (Stealth Biotech)

Monday 21 Sept 2026, 16:15 → 17:30 Europe/Zurich

40/S2-D01 - Salle Dirac (CERN)

Introduction to Differentiable Scientific Computing in JAX + Equinox

Automatic differentiation is transforming scientific computing, enabling gradient-based approaches to optimization. It powers modern artificial intelligence, and now also transforms the world of scientific computing, enabling end-to-end workflows bridging both worlds. 

This talk introduces how JAX's composable transformations and the Equinox ecosystem make this practical: Equinox provides module system built on JAX's functional paradigm with Pytorch-like syntax, while libraries such as Diffrax for differential equations, Lineax for linear solvers, and Optimistix for nonlinear optimization and root-finding, offer performant numerics out of the box. We'll walk through key concepts and real examples showing how this ecosystem enables end-to-end differentiable scientific workflows.

---

CAUTION: max. number of persons in Salle Dirac is 47!

 

equinox_introduction.pdf