Faster than TensorFlow and PyTorch, JAX is the new open source Google framework 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).

This course will introduce the basics of JAX functioning—its strengths and its constraints.