flâneur

Just-in-time compilation — JAX documentation

jax.readthedocs.io · 2,180 words · saved by 1 readers

In this section, we will further explore how JAX works, and how we can make it performant. We will discuss the jax.jit() transformation, which will perform Just In Time (JIT) compilation of a JAX Python function so it can be executed efficiently in XLA. In the previous section, we discussed that JAX allows us to transform Python functions. JAX accomplishes this by reducing each function into a sequence of primitive operations, each representing one fundamental unit of computation. One way to see the sequence of primitives behind a function is using jax.make_jaxpr(): The Understanding Jaxprs section of the documentation provides more information on the meaning of the above output. Importantly, notice that the jaxpr does not capture the side-effect present in the function: there is nothing in it corresponding to global_list.append(x). This is a feature, not a bug: JAX transformations are designed to understand side-effect-free (a.k.a. functionally pure) code. If pure function and side-ef

Just-in-time compilation # In this section, we will further explore how JAX works, and how we can make it performant. We will discuss the jax.jit() transformation, which will perform Just In Time (JIT) compilation of a JAX Python function so it can be executed efficiently in XLA. How JAX transformations work # In the previous section, we discussed that JAX allows us to transform Python functions. JAX accomplishes this by reducing each function into a sequence of primitive operations, each representing one fundamental unit of computation. One way to see the sequence of primitives behind a funct

related reading