The moment when JIT compilation is done

[GoktugGuvercin]

Hello JAX Community;

While studying just-in-time compilation in JAX, I am confused about one thing, and I am here to ask it. In official JAX documentation, selu activation is given as an example function to be jit-compiled, and then its jit compilation is explained with these sentences:

import jax
import jax.numpy as jnp

def selu(x, alpha=1.67, lambda_=1.05):
  return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

x = jnp.arange(1000000)
%timeit selu(x).block_until_ready()

selu_jit = jax.jit(selu)

# Warm up
selu_jit(x).block_until_ready()

%timeit selu_jit(x).block_until_ready()

We ran selu_jit once on x. This is where JAX does its tracing – it needs to have some inputs to wrap in tracers, after all. The jaxpr is then compiled using XLA into very efficient code optimized for your GPU or TPU. Subsequent calls to selu_jit will now use that code, skipping our old Python implementation entirely.

What I want to ask is actually first call. When selu_jit() is called first time, is function execution carried out over unoptimized, standard form or its optimized version ?

I think that first call should be done over unoptimized (without operation fusion of XLA) version, since during the first call, tracing and instantiation of jaxpr are performed, so before these ones, XLA cannot compile and optimize selu I guess.

[jakevdp]

The first time selu_jit() is called, it is traced & compiled using abstract values representing the input arrays, and then once compilation is finished, the actual arrays are passed to the compiled function for execution. The tracing and compilation is done with abstract values representing the input array, not with the input array itself.

As a side note, very early versions of JAX did this differently, and the first function execution was unoptimized, with only subsequent executions being compiled. But JAX has not worked that way for a long time.

[GoktugGuvercin]

If first execution is performed over compiled and optimized function after the tracing and then compilation over abstract values, why does first execution contain side effects ?

In the documentation, it says that each argument of python function is wrapped by a tracer object, and then each of those tracers take a record of all operations defined on its wrapped argument. Composition of those records eventually constructs jaxpr instance. During tracing, any side-effect is not included into jaxpr context. This jaxpr is used to be compiled by XLA for optimized code. If first execution (first call to selu_jit()) is performed over compiled function, then it should be side-effect free, but first execution runs the side effects. In the code below, counter is incremented by 1.

import jax
import jax.numpy as jnp

counter = 0

def selu(x, alpha=1.67, lambda_=1.05):
  global counter
  counter = counter + 1
  return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

x = jnp.arange(1000000)
selu_jit = jax.jit(selu)

# Warm up
# this changed the value for counter, but I think it should not change if first call is done over compiled function
selu_jit(x).block_until_ready()
print("counter = ", counter)

%timeit selu_jit(x).block_until_ready()

Did I misunderstand compilation and side-effect concept ?

[jakevdp]

Side-effects happen during tracing, not execution. So the first time you call the function, it gets traced, during which any side-effects happen, then compiled by XLA, and then executed without side-effects. The second time you call the function on inputs with the same shape and dtype, the cached version of the compiled function is used, so there is no tracing, side-effects, or compilation; just side-effect-free execution.

[GoktugGuvercin]

Waow, I completely get it right now. Thank you so much for your clarifications.