Need to speed up a large neural-network computation, climate model, or even an image filter? JAX compilation transforms Python computations into executable form so repeated or massive workloads can run faster. The main entry point is `jax.jit`, demonstrated with a two-branch SELU activation function. section 5.1
JAX first converts a function into Jaxpr, its explicitly typed intermediate representation. Transformations operate on this functional representation before the computation is sent to XLA. `jax.make_jaxpr()` exposes operations, shapes, types, constants, and intermediate variables. section 5.2.1
Compilation works best when the traced computation has a form JAX can represent:
- Use `jax.make_jaxpr()` to inspect how array operations become primitives such as `add`, `mul`, broadcasting, and `reduce_sum`. section 5.2.1
- Expect ordinary Python loops with fixed bounds to be expanded during tracing, while input-dependent bounds can raise `TracerIntegerConversionError`. section 5.2.1
- Replace input-dependent control flow with `jax.lax.fori_loop` or `jax.lax.cond`; alternatively, mark values static when specializing compilation to concrete inputs. section 5.2.1
Inspecting a computation
For `f1(x, y, z) = jnp.sum(x + y * z)`, calling `jax.make_jaxpr(f1)(x, y, z)` reveals broadcasts, multiplication, addition, and reduction, including each array's shape and type. JIT-compiling the equivalent structured-control-flow functions preserves their results while producing executable JAX arrays. section 5.2.1 section 5.2.1
Listing 5.10 Using jax.make_jaxpr()
def f1(x, y, z): #1
return jnp.sum(x + y * z)
x = jnp.array([1.0, 1.0, 1.0]) #2
y = jnp.ones((3,3))*2.0
z = jnp.array([2.0, 1.0, 0.0]).T
jax.make_jaxpr(f1)(x,y,z) #3
>>> { lambda ; a:f32[3] b:f32[3,3] c:f32[3]. let
>>> d:f32[1,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 3)] c
>>> e:f32[3,3] = mul b d
>>> f:f32[1,3] = broadcast_in_dim[broadcast_dimensions=(1,) shape=(1, 3)] a
>>> g:f32[3,3] = add f e
>>> h:f32[] = reduce_sum[axes=(0, 1)] g
>>> in (h,) }
Inspect the complete Jaxpr generated for a small array computation. section 5.2.1
