Modeling and Simulation Group Meeting

---

JAX: A platform for high-performance numerical computing with applications in ML research and beyond

Speaker: Dan Foreman-Mackey, Google Deepmind

Location: Warren Weaver Hall 202

Date: Thursday, October 10, 2024, 12:30 p.m.

Synopsis:

In this interactive and informal seminar, I will introduce JAX (https://github.com/jax-ml/jax), an open source project that enables differentiable, parallelizable, and hardware accelerated numerical computing. JAX is widely used for machine learning research, but it is also a useful tool for other computing tasks, including the area where I have the most experience: probabilistic programming. After demonstrating the core JAX programming model, I will share "tinygp" (a tiny Gaussian Process package that I maintain, https://github.com/dfm/tinygp) as a case study for how the community can build on top of JAX to build performant and user-friendly Python libraries. This seminar will include some live coding, and plenty of time for discussion.