5 Compiling your code

In chapter 1, we compared the performance of a simple JAX function on a CPU and GPU, with and without JIT. In chapter 2, we used JIT to compile two functions in a training loop for a simple neural network. So you basically know what JIT does. It compiles your function for a target hardware platform and makes it faster.

Beginning in chapter 4, we started learning JAX transformations (remember, JAX is about composable function transformations!). That chapter taught us about autodiff and the grad() transformation. This chapter will discuss compilation and the corresponding jit() transformation. In the following chapters, we will learn more transformations related to auto-vectorization and parallelization.

In the landscape of numerical computing, JAX emerged as a compelling framework rooted in the foundational strength of Google’s XLA compiler. The XLA compiler is not just another tool in the computational toolkit; it is specifically tailored to produce efficient code for high-performance computing tasks. Think of it as an architect meticulously designing blueprints to build structures optimized for their purpose.

5.1 Using compilation

5.1.1 Using JIT compilation

5.1.2 Pure functions and compilation process

5.2 JIT internals

5.2.1 Jaxpr, an intermediate representation for JAX programs

5.2.2 XLA

5.2.3 Using AOT compilation

5.3 JIT limitations

5.3.1 Pure and impure functions

5.3.2 Exact numerics

5.3.3 Conditioning on input parameter values

5.3.4 Slow compilation

5.3.5 Class methods

5.3.6 Simple functions

Exercise 5.1

Summary