JAX: a framework for high-performance array computing


Marie-Hélène Burle

JAX is an open source framework by Google for high-performance array computations and differentiation. With a NumPy-like API, it uses asynchronous dispatch, just-in-time compilation, and the XLA compiler for linear algebra to run on accelerators (GPUs and TPUs). It is both very fast and flexible. This however comes at the cost of some constraints: JAX requires pure functional programming code.

This webinar will introduce JAX’s principles and functioning.

Coming up in spring 2024.