XLA runner (to support JAX)

[VivekPanyam]

Background

XLA is an ML "compiler for GPUs, CPUs, and ML accelerators."

Carton support for XLA would primarily be used to provide JAX support, but in theory it could also support some PyTorch and TensorFlow models.

Here's a guide on how to export a JAX model from Python and run it from C++ using XLA: jax-ml/jax#5337 (comment).

This is an example of the above in the JAX codebase.

I've explored doing this in the past (outside of Carton), but there weren't XLA prebuilt binaries available and it required building from source in the TensorFlow repo. Now, with OpenXLA and prebuilt binaries, this is a lot easier.

@LaurentMazare created rust bindings to XLA that include a straightforward example of loading the HLO IR generated by the JAX export code. That should make it fairly easy to prototype an integration with Carton if anyone is interested in doing so.

Implementation

Concretely, this could be implemented as follows:

• A Python utility in the Python bindings that accepts a function, does some of what jax_to_ir.py does and calls jax.xla_computation
• A new Carton runner crate that uses the XLA crate linked above to load and run the HLO

Recommended reading:

• https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html
• https://github.com/VivekPanyam/carton/blob/main/CONTRIBUTING.md#adding-support-for-a-new-ml-framework