Skip to main content

Developing software tools for accelerated and differentiable scientific computing using JAX

H.1308 (Rolin) | Day 2 | 13:20 - 13:30 | Speakers: Matt Graham

Developing software tools for accelerated and differentiable scientific computing using JAX
A picture of a devroom at FOSDEM 2024
Open in browser
Get involved in the conversation!Join the chat

Notes

Abstract

JAX is an open-source Python package for high-performance numerical computing. It provides a familiar NumPy style interface but with the advantages of allowing computations to be dispatched to accelerator devices such as graphics and tensor processing units, and supporting transformations to automatically differentiate, vectorize and just-in-time compile functions. While extensively used in machine learning applications, JAX's design also makes it ideal for scientific computing tasks such as simulating numerical models and fitting them to data.

This lightning talk will introduce JAX's interface and computation model, and some of its key function transformations. I will also briefly introduce the Python Array API standard and explain how it can be used to write portable code which works across JAX, NumPy and other array backends.

Speakers

Matt Graham

Matt Graham is a research data scientist at the UCL Centre for Advanced Research Computing. His work focuses on designing, implementing and applying algorithms for efficiently calibrating probabilistic models to data, with a particular focus on the use of differentiable programming in this context. He has extensive experience in software development for scientific computing, especially within the scientific Python ecosystem. He has contributed to numerous open-source software projects, including Theano, Mici, S2FFT, Thanzi la Onse model and ParticleDA.jl.


Notice: The placeholder video image is licensed under CC BY-SA 4.0. The original image can be found hereChanges made to the image are: Cropped the image to a new ratio, part of the image was cut off.