Skip to content

← JAX Profiling & Debugging step 11 of 12

Hard Primitives

AOT via jit.lower().compile()

Jump to the task

Compile before you call

Ordinary jax.jit combines several stages behind a callable: tracing, lowering to a compiler representation, compiling, then executing. You can request the preparation stages explicitly with .lower(...).compile() and keep the returned executable for later calls.

This is ahead-of-time compilation relative to execution. It can move compilation out of a request-handling path, but does not by itself create a portable deployment file or guarantee identical first-call latency. Transfers, runtime initialization and other work may still take time.

Worked example

import jax
import jax.numpy as jnp

example = jnp.ones(4, dtype=jnp.float32)
jitted_sum = jax.jit(jnp.sum)
lowered = jitted_sum.lower(example)
compiled = lowered.compile()

first = compiled(example)             # 4.0
second = compiled(jnp.arange(4.0))     # 6.0

The executable accepts different values with the same input signature. It was prepared for shape (4,) and dtype float32, not for the particular four ones in example. Calling it with a length-three array is a signature mismatch; it does not automatically compile a new version.

Problem

Implement aot_compiled_sum(x):

  1. Wrap jnp.sum with jax.jit.
  2. Call .lower(x).compile() on the wrapped function.
  3. Call the resulting compiled function on x and return its scalar sum.

x is a one-dimensional float32 JAX array. Empty input has sum zero. Do not return a lowered object or the compiled callable itself.

This small exercise deliberately performs preparation and execution in one function so you can practice the API. For repeated real-world calls, prepare once outside the repeated path and retain the compiled object, as the example does. These numerical tests check the sum, not compilation reuse or performance.

Understand the two preparation steps

lower(x) returns a lowered computation object, not just a text string. You can inspect it with .as_text() or .compiler_ir(). Calling .compile() on it produces an executable specialized to the input signature and compilation environment.

You can supply jax.ShapeDtypeStruct((4,), jnp.float32) to lower instead of allocating example values. The shape and dtype describe the input; you still pass actual arrays when executing the compiled result.

A compiled callable is not interchangeable with a normal JAX function: you cannot directly apply transformations such as jax.grad or jax.vmap to it. Apply those transformations before lowering and compiling.

Multiple shapes and deployment

A regular jitted function can create specializations for new shapes. An individual compiled callable accepts the signature it was compiled for. Shape-polymorphic export can reuse tracing and lowering across shapes, but it is not a promise of one device binary for every shape; concrete shapes can still need separate compilation.

Exporting or serializing a computation is a separate workflow from keeping this in-process compiled object. Backend compatibility and runtime setup still matter. Use AOT here to control when preparation occurs, not as a claim of portable or automatically reproducible binaries.

The diagram illustrates ordinary JIT specialization for lengths 32 and 64. The x32/x64 labels are lengths, not dtypes. Unlike that reusable JIT wrapper, a single compiled callable does not add a new specialization on demand.

Loading visualization…