Just-in-time compilation

Just-in-time compilation#

jax.jit() is the transformation that makes JAX code fast: it compiles a Python function into a single optimized computation, fused and specialized for your hardware. This page is about what compilation buys you, when it happens (and re-happens), and the levers you control: static arguments, caching, in-place updates with refs and buffer donation, and asynchronous dispatch.

What compilation buys you#

Without jit, JAX executes operations (think jnp function calls) one at a time, dispatching each to the device as it’s encountered. That’s flexible in that it allows arbitrary Python control flow, but it limits performance in two ways: each dispatch carries Python and runtime overhead, and a compiler never gets to see the whole computation at once, so it can’t optimize across operations.

jax.jit addresses both limitations. It traces the function to a jaxpr (How transformations work: tracing) and hands the whole jaxpr to XLA, which compiles it into a single optimized executable, fusing operations, eliminating temporary arrays, and specializing for your CPU, GPU, or TPU. Consider a scaled exponential linear unit (SELU), an operation commonly used in deep learning:

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()

6.89 ms ± 201 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

selu_jit = jax.jit(selu)

# Warm up: the first call traces and compiles.
selu_jit(x).block_until_ready()

%timeit selu_jit(x).block_until_ready()

657 μs ± 1.59 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Here’s what just happened:

  1. We defined selu_jit as the compiled version of selu.

  2. We called selu_jit once on x. This is where the tracing and XLA compilation happen. (We do this warm-up call outside the timing loop so that we measure execution, not compilation.)

  3. We timed the compiled version. Subsequent calls with the same input types never run the Python function selu at all; they go straight to the compiled executable.

(Note the use of block_until_ready(): because of JAX’s asynchronous dispatch, covered below, timing without it would measure only how long it takes to launch the work.)

What the compiled function captures#

A jaxpr records the function as executed on the JAX types it was traced with. Python-level control flow and side effects are resolved during tracing. If a function branches on a static property like rank, the jaxpr knows only the branch that was taken:

def log2_if_rank_2(x):
  if x.ndim == 2:
    return jnp.log(x) / jnp.log(2.0)
  else:
    return x

jax.jit(log2_if_rank_2).trace(jnp.array([1, 2, 3])).jaxpr

{ lambda ; a:i32[3]. let  in (a,) }

For a rank-1 input, the compiled function is just the identity. That’s fine, because the cached executable is only reused for inputs of the same JAX type; a rank-2 input triggers a fresh trace that takes the other branch.

Impure functions are dangerous under jit for exactly this reason: side effects happen once, at trace time, and are then absent from the cached executable, so they might appear to work on the first call and silently vanish afterwards. JAX often can’t detect the impurity. For debug printing that survives compilation, use jax.debug.print(); for general side effects at a performance cost, see External callbacks. To check for leaked tracers from side effects, use jax.check_tracer_leaks().

Why can’t we just jit everything?#

Recall from Transformations: grad and vmap that traced code can’t always specialize on data values — and jit is the extreme case: its tracers carry no values at all, only JAX types. Value-dependent Python control flow therefore fails:

# Condition on the value of x.

def f(x):
  if x > 0:
    return x
  else:
    return 2 * x

jax.jit(f)(10)  # Raises an error

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function f at /tmp/ipykernel_1614/3856269748.py:3 for jit. This concrete value was not available in Python because it depends on the value of the argument x.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError

# While loop conditioned on x and n.

def g(x, n):
  i = 0
  while i < n:
    i += 1
  return x + i

jax.jit(g)(10, 20)  # Raises an error

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function g at /tmp/ipykernel_1614/722961019.py:3 for jit. This concrete value was not available in Python because it depends on the value of the argument n.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError

The problem in both cases is that we tried to condition the program’s trace-time flow on runtime values. Traced values inside jit, like x and n here, can only affect control flow through their JAX types, not their values. For much more on this, including the structured control-flow operations jax.lax.cond and jax.lax.scan that express data-dependent control flow inside compiled code, see Control flow and logical operators with jit.

One pragmatic option is to jit only part of a function. If the expensive work is inside the loop body, compile just the body and let Python drive the loop:

@jax.jit
def loop_body(prev_i):
  return prev_i + 1

def g_inner_jitted(x, n):
  i = 0
  while i < n:
    i = loop_body(i)
  return x + i

g_inner_jitted(10, 20)

Array(30, dtype=int32, weak_type=True)

This gives up cross-operation optimization for the loop itself, but keeps each iteration compiled. (If you do this, read JIT and caching below to avoid accidentally recompiling every iteration.)

Marking arguments as static#

To branch on the value of an argument, one option is to treat that argument as static: a regular Python value fixed at trace time, rather than a traced array. The static_argnums and static_argnames parameters to jax.jit do this:

f_jit_correct = jax.jit(f, static_argnums=0)
print(f_jit_correct(10))

10

g_jit_correct = jax.jit(g, static_argnames=['n'])
print(g_jit_correct(10, 20))

30

To specify static arguments when using jit as a decorator, use the decorator-factory pattern:

@jax.jit(static_argnames=['n'])
def g_jit_decorated(x, n):
  i = 0
  while i < n:
    i += 1
  return x + i

print(g_jit_decorated(10, 20))

30

The compiled executable is specialized to the value of each static argument, so JAX recompiles whenever a static argument takes a value it hasn’t seen. Marking an argument static is only a good strategy when the function will see a limited set of values for it. Static argument values must be hashable, since they become part of the cache key.

JIT and caching#

The first call to a jitted function performs tracing and compilation. Subsequent applications of the function to the same input types will instead use cached compilations. Understanding how the compilation cache works is important: several common patterns silently defeat it.

Suppose we define f = jax.jit(g). When f is first invoked, it gets compiled and the resulting executable is cached. The cache key includes:

  • the JAX types of the arguments — new shapes or dtypes trigger a fresh trace and compile;

  • the values of any static arguments; and

  • the object identity of the function g itself.

That last point has a practical consequence: avoid calling jax.jit on temporary functions defined inside loops or other inner scopes. The cache relies on the object identity of the function, so freshly-created lambdas and partial objects — even ones wrapping the same underlying code — look like new functions every time, and recompile every time:

from functools import partial

def unjitted_loop_body(prev_i):
  return prev_i + 1

def g_inner_jitted_partial(x, n):
  i = 0
  while i < n:
    # Don't do this! A fresh `partial` has a different hash each time,
    # so this recompiles on every iteration.
    i = jax.jit(partial(unjitted_loop_body))(i)
  return x + i

def g_inner_jitted_lambda(x, n):
  i = 0
  while i < n:
    # Don't do this either! Same problem: a fresh lambda every iteration.
    i = jax.jit(lambda x: unjitted_loop_body(x))(i)
  return x + i

def g_inner_jitted_normal(x, n):
  i = 0
  while i < n:
    # This is OK: `jax.jit` sees the same function object each time,
    # and finds the cached executable.
    i = jax.jit(unjitted_loop_body)(i)
  return x + i

print("jit called in a loop with partials:")
%timeit g_inner_jitted_partial(10, 20).block_until_ready()

print("jit called in a loop with lambdas:")
%timeit g_inner_jitted_lambda(10, 20).block_until_ready()

print("jit called in a loop with caching:")
%timeit g_inner_jitted_normal(10, 20).block_until_ready()

jit called in a loop with partials:
396 ms ± 6.17 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
jit called in a loop with lambdas:
410 ms ± 10.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
jit called in a loop with caching:
3.07 ms ± 17 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)

When a program recompiles more than you expect, try setting the jax_explain_cache_misses config option to make JAX print an explanation for every tracing cache miss:

jax.config.update("jax_explain_cache_misses", True)

@jax.jit
def double(x):
  return 2 * x

_ = double(jnp.arange(3))    # cache miss: never seen this function before
_ = double(jnp.arange(3))    # cache hit: silent
_ = double(jnp.arange(3.0))  # cache miss: x's JAX type changed

TRACING CACHE MISS at /tmp/ipykernel_1614/401824424.py:7:4 (<module>):
  never seen function:
    double id=126450749752800 defined at /tmp/ipykernel_1614/401824424.py:3
TRACING CACHE MISS at /tmp/ipykernel_1614/401824424.py:9:4 (<module>):
  for double defined at /tmp/ipykernel_1614/401824424.py:3
  all previously seen cache keys differ. For the closest previous key:
  different input types:
  * at x, now f32[3] and before i32[3]

jax.config.update("jax_explain_cache_misses", False)

For a lighter-weight signal, jax_log_compiles logs one line per compilation, without the diagnosis. When tracing or compilation time itself becomes the problem, Debugging slow JAX tracing and XLA compilation is a field guide to diagnosing it.

Another classic source of endless recompilation is data whose shape genuinely varies from call to call: ragged sequence lengths, a partial final batch at the end of an epoch. Since the cache is keyed on JAX types, every new shape is a fresh compile. The standard fix is to pad variable-shaped data up to a small, fixed set of bucket shapes, trading a little wasted compute (and a masking jnp.where or two) for a bounded number of compilations.

The persistent compilation cache (Persistent compilation cache) keeps compiled executables across process restarts; in the other direction, jax.clear_caches() drops every cached trace and executable.

Using jit with methods#

Most jax.jit examples decorate stand-alone functions, but sooner or later you’ll want to jit a method:

class CustomClass:
  def __init__(self, x: jax.Array, mul: bool):
    self.x = x
    self.mul = mul

  @jax.jit  # <---- How to do this correctly?
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

c = CustomClass(2, True)
c.calc(3)

TypeError: Error interpreting argument to <function CustomClass.calc at 0x73019c055080> as an abstract array. The problematic value is of type <class '__main__.CustomClass'> and was passed to the function at path self.
This typically means that a jit-wrapped function was called with a non-array argument, and this argument was not marked as static using the static_argnums or static_argnames parameters of jax.jit.

The problem is that the first argument to calc is self, of type CustomClass, and JAX doesn’t know how to handle that type. There are three basic strategies, in increasing order of power.

Strategy 1: a jitted helper function#

The most straightforward approach: keep jit on a stand-alone helper that takes only JAX-compatible arguments, and have the method call it:

class CustomClass:
  def __init__(self, x: jax.Array, mul: bool):
    self.x = x
    self.mul = mul

  def calc(self, y):
    return _calc(self.mul, self.x, y)

@jax.jit(static_argnums=0)
def _calc(mul, x, y):
  if mul:
    return x * y
  return y

c = CustomClass(2, True)
print(c.calc(3))

6

This is simple and explicit, and JAX never has to learn about CustomClass at all. The cost is that the method’s logic lives outside the class.

Strategy 2: marking self as static#

Another common pattern is to mark the self argument as static. Done naively, it misbehaves:

class CustomClass:
  def __init__(self, x: jax.Array, mul: bool):
    self.x = x
    self.mul = mul

  # WARNING: broken, as we'll see below. Don't copy & paste!
  @jax.jit(static_argnums=0)
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

c = CustomClass(2, True)
print(c.calc(3))

6

This runs, but it has a flaw. Static arguments become cache keys, so JAX relies on their hash and equality. The default __hash__ for a custom object is its object ID, which doesn’t change when the object mutates — so mutating the object silently serves stale compiled code:

c.mul = False
print(c.calc(3))  # Should print 3... but the cache doesn't know mul changed

6

You can partially address this by defining __hash__ and __eq__ in terms of the object’s contents, so that differing objects actually miss the cache:

class CustomClass:
  def __init__(self, x: jax.Array, mul: bool):
    self.x = x
    self.mul = mul

  @jax.jit(static_argnums=0)
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

  def __hash__(self):
    return hash((self.x, self.mul))

  def __eq__(self, other):
    return (isinstance(other, CustomClass) and
            (self.x, self.mul) == (other.x, other.mul))

This works with jit and other transformations so long as you never mutate the object: mutating a value that’s in use as a hash key leads to subtle problems — which is exactly why Python’s mutable containers (dict, list) don’t define __hash__. If your class mutates its attributes, it isn’t really static, and there’s a better option.

Strategy 3: registering the class as a pytree#

The most flexible approach is to register the type as a custom pytree node (Custom pytree nodes), saying explicitly which parts are dynamic data and which are static metadata:

from jax import tree_util

class CustomClass:
  def __init__(self, x: jax.Array, mul: bool):
    self.x = x
    self.mul = mul

  @jax.jit
  def calc(self, y):
    if self.mul:
      return self.x * y
    return y

  def _tree_flatten(self):
    children = (self.x,)          # arrays / dynamic values
    aux_data = {'mul': self.mul}  # static values
    return (children, aux_data)

  @classmethod
  def _tree_unflatten(cls, aux_data, children):
    return cls(*children, **aux_data)

tree_util.register_pytree_node(CustomClass,
                               CustomClass._tree_flatten,
                               CustomClass._tree_unflatten)

Now self flows through jit like any other pytree argument: the arrays in children are traced, the aux_data is static, plain @jax.jit works with no static_argnums, and the problems above solve themselves:

c = CustomClass(2, True)
print(c.calc(3))

c.mul = False                        # mutation is detected
print(c.calc(3))

c = CustomClass(jnp.array(2), True)  # non-hashable x is supported
print(c.calc(3))

6
3
6

If your class is a dataclass, jax.tree_util.register_dataclass() gets you the same result with much less ceremony; see Custom pytree nodes.

In-place updates with refs#

A compiled function’s inputs and outputs are immutable arrays, so at the jit boundary every output gets a fresh buffer, even when you’re conceptually just updating an input, as in a training step that maps parameters to new parameters. Inside the compiled computation XLA reuses memory aggressively, but it can’t reuse an input’s buffer for an output unless you tell it something extra. JAX has two ways to say it: pass the data as a ref, or donate the input buffer. Refs express the intent directly, so we cover them first; donation, next, is the longer-standing mechanism you’ll meet throughout existing JAX code.

A ref (Refs: mutable arrays) passed into a jitted function represents mutable memory: writes through it happen in place, and no fresh buffer is allocated. We can watch the buffer address stay fixed across the call:

@jax.jit
def sin_inplace(x_ref):
  x_ref[...] = jnp.sin(x_ref[...])

x_ref = jax.new_ref(jnp.arange(3.0))
print(x_ref.unsafe_buffer_pointer(), x_ref)
sin_inplace(x_ref)
print(x_ref.unsafe_buffer_pointer(), x_ref)

104040855109504 Ref([0., 1., 2.], dtype=float32)
104040855109504 Ref([0.        , 0.84147096, 0.9092974 ], dtype=float32)

sin_inplace updates the buffer backing x_ref, so its address stays the same. In general, under a jit you should expect refs to point to fixed buffer addresses, and indexed updates to be performed in place.

Note

Temporary caveat: dispatch from Python to impure jit-compiled functions that take ref inputs is currently slower than dispatch to pure jit-compiled functions, since it takes a less optimized path.

Working with refs does mean writing the function differently: mutation through indexed writes, rather than functional updates. For code already written functionally (which is most existing JAX code, and the convention across the ecosystem), buffer donation can achieve the in-place effect without changing the function.

Buffer donation#

When JAX executes a computation, it uses device buffers for all inputs and outputs. If you know an input won’t be needed after the computation, and it matches the shape and element type of an output, you can donate the input buffer to hold that output, reducing peak memory by the size of the donated buffer. The canonical case is a functional update, where the new state replaces the old:

params, state = jax.jit(update_fn, donate_argnums=(0, 1))(params, state)

Think of donation as a memory-efficient functional update on immutable arrays. Within a compiled computation, XLA reuses buffers automatically; at the jit boundary, though, JAX must assume you might still hold a reference to the input unless you promise otherwise with donate_argnums (or donate_argnames):

def add(x, y):
  return x + y

x = jax.device_put(np.ones((2, 3)))
y = jax.device_put(np.ones((2, 3)))
# Execute `add` with donation of the buffer for `y`. The result has
# the same shape and type as `y`, so it will share its buffer.
z = jax.jit(add, donate_argnums=(1,))(x, y)

Donation comes with rules and sharp edges:

  • Donated means gone. After the call, the donated input is invalid, and using it is an error:

    z = jax.jit(add, donate_argnums=(1,))(x, y)
w = y + 1  # Reuses `y`, whose buffer was donated above
# >> RuntimeError: Invalid argument: CopyToHostAsync() called on invalid buffer

  • Keyword arguments aren’t donated by donate_argnums. This code donates nothing:

    params, state = jax.jit(update_fn, donate_argnums=(0, 1))(params=params, state=state)

  • Pytree arguments donate all their buffers. Donating an argument that’s a pytree donates every array in it:

    def add_ones(xs: list[jax.Array]):
  return [x + 1 for x in xs]

xs = [jax.device_put(np.ones((2, 3))), jax.device_put(np.ones((3, 4)))]
# Donates the buffers of both arrays in `xs`.
z = jax.jit(add_ones, donate_argnums=0)(xs)

  • Unusable donations are dropped, with a warning. If there are more donated buffers than outputs to hold them, or no output matches a donated buffer’s shape and element type, you’ll see UserWarning: Some donated buffers were not usable.

As the sharp edges above suggest, donation is a promise layered on top of immutable semantics, and it’s possible to get subtly wrong. The promise is checked only partially, and only at runtime. When you have the freedom to restructure, refs (In-place updates with refs above) express the same intent with less to misuse. Donation remains the workhorse for functionally-written code, which is to say most JAX code today.

Asynchronous dispatch#

One more piece of the performance model, relevant with and without jit: JAX does not wait for a computation to finish before returning control to Python. Consider:

import numpy as np
from jax import random

x = random.uniform(random.key(0), (1000, 1000))
jnp.dot(x, x) + 3.0

Array(shape=(1000, 1000), dtype=float32)

When jnp.dot(x, x) executes, JAX returns immediately with a jax.Array that is really a future: a value that will be produced on the device but isn’t necessarily available yet. We can inspect its JAX type, and pass it into further JAX computations, without waiting; only when we actually look at the value from the host (printing it, converting it to a NumPy array) does Python block until the computation is done. (The cell above shows values only because rendering the result forced the wait.)

This is why asynchronous dispatch is valuable: Python “runs ahead” of the device, enqueueing work and staying off the critical path.

Asynchronous operations make timing a bit trickier. Time an operation naively and you measure only the dispatch:

%time jnp.dot(x, x)

CPU times: user 3.61 ms, sys: 0 ns, total: 3.61 ms
Wall time: 1.86 ms

Array(shape=(1000, 1000), dtype=float32)

A fraction of a millisecond would be a suspiciously good time for a 1000Ă—1000 matrix multiplication! To measure the actual computation, force completion with block_until_ready():

%time jnp.dot(x, x).block_until_ready()

CPU times: user 56.9 ms, sys: 691 ÎĽs, total: 57.6 ms
Wall time: 29.4 ms

Array(shape=(1000, 1000), dtype=float32)

Blocking without transferring the result to the host (as block_until_ready does) is usually faster than forcing a transfer with np.asarray(...), and is the right tool for microbenchmarks. Benchmarking JAX code has a few more pitfalls of this flavor — Benchmarking and profiling starts there.

Where to next#