Need predictable startup time instead of paying compilation cost on the first call? Ahead-of-time (AOT) compilation moves compilation before execution, so the resulting function runs without a JIT warm-up. section 5.2.3 section 5.1.1
JAX first stages a Python function for a fixed input type signature, lowers the staged representation to StableHLO through MLIR and XLA, and compiles device-specific code for CPU, GPU, or TPU. The AOT API controls the lowering and compilation stages before invocation. section 5.2.3 section 5.5
Compile SELU before calling it
A JIT-wrapped function compiles when first invoked, while an AOT function compiles directly: selu_jit = jax.jit(selu) selu_aot = jax.jit(selu).lower(1.0).compile() The side effect inside selu runs during the direct AOT compilation, but during the first call to selu_jit. Both functions then produce the same result for the compiled signature. section 5.2.3
Keep these constraints in mind:
- Use AOT when compilation timing matters; use JIT when compilation on first use and later reuse better fit the call pattern. section 5.2.3 section 5.1.1
