Saving JIT'ed functions, and Ahead-of-Time Compilation

[jusevitch]

I'm new to JAX, and I had a couple questions about its compilation capabilities:

1. Is there a way to save JIT'ed functions so that they persist across REPL / Python sessions? Or do the JIT'ed functions disappear when you close the Python interpreter, similar to what happens in the Julia programming language?
2. Are there any capabilities for ahead-of-time (AOT) compilation using JAX? Numba has some capabilities for this (see these docs), and I was wondering if JAX can do the same (or if it's planned for future development).

[YouJiacheng]

Maybe check https://github.com/google/jax/blob/main/jax/_src/dispatch.py#L555.
Seems that currently only support TPU.

[froystig]

1. Is there a way to save JIT'ed functions so that they persist across REPL / Python sessions? Or do the JIT'ed functions disappear when you close the Python interpreter, similar to what happens in the Julia programming language?

There's an experimental persistent cache that works on TPU only today. See jax.experimental.compilation_cache. Documentation is still on the way.

2. Are there any capabilities for ahead-of-time (AOT) compilation using JAX? Numba has some capabilities for this (see these docs), and I was wondering if JAX can do the same (or if it's planned for future development).

#7733 is in development, usable today but not yet officially stable. This effectively offers "AOT" within a running Python interpreter process, but doesn't persist across process restarts. Another way to describe it is as "an API for explicit control of lowering and compilation."