ahead-of-time lowering and compilation

Following the experiment in #6034, plus other thoughts and motivation accrued since, we should support a notion of lowering and compilation “ahead of time,” i.e. upon explicit request. Some things to note:

• Lowering means producing XLA HLO. Compiling means running it through the XLA compiler
• Compiling produces optimized HLO and renders a computation ready for execution.
• Lowering and compilation happen for a fixed type signature.
• A function compiled ahead of time and can be called with arguments of the fixed type signature only.
• A function compiled ahead of time will intentionally fail to transform, e.g. if called under jvp or vmap, since it will be compiled for a particular type signature, and transformations modify signatures.
• We will want this for every compiling transformation: jit, pjit, pmap, xmap.
• Although a function compiled ahead of time is bound to correspond to some XLA executable, it’s out of scope to provide supported access to, or any serialization of, that executable.

[froystig]

With documentation now in place (#12204) and plenty of usage already, it’s time to declare this complete! We will return with any thoughts and work on cross-AOT when the time comes for that.

[froystig]

The API and implementation are essentially in place, and already in use. I’m working on documentation.