Skip to content

Compilation¤

End-to-end JAX training loops are compiled as one big program, which can take a while. Jaxnasium provides two convenience helpers around JAX's own facilities: one to compile ahead of time (with progress reporting), and one to persist compilation results across processes.

import jax
import jaxnasium as jym
from jaxnasium.algorithms import PPO

jym.enable_compilation_cache()

env = jym.make("CartPole-v1")
train = jym.precompile(PPO().train, jax.random.PRNGKey(0), env)
agent, metrics = train()

jaxnasium.precompile(fn: Callable[..., Any], *args: Any, **kwargs: Any) -> CompiledFunction ¤

Compile fn(*args, **kwargs) ahead of time, with some progress to stderr to keep you entertained while compiling; especially since compiling the full training loop can take a while.

https://docs.jax.dev/en/latest/aot.html

Additionally, the returned function stores the arguments it was compiled for, so you can call it directly with no arguments.

train = jym.precompile(ppo.train, key, env)
# Compiling `train` ...
#   ✓ traced and lowered in 6.3s
#   ✓ compiled with XLA in 12.8s
# ✓ `train` ready in 19.1s
agent, metrics = train()

jaxnasium.enable_compilation_cache(path: str | os.PathLike | None = None, *, min_compile_time_secs: float = 0.01, min_entry_size_bytes: int = 0) -> Path ¤

Convenience function to enable JAX's persistent compilation cache.

Default path is $CACHE_HOME_DIR/jaxnasium/jit-cache or $XDG_CACHE_HOME/jaxnasium/jit-cache or ~/.cache/jaxnasium/jit-cache.

https://docs.jax.dev/en/latest/persistent_compilation_cache.html