I think this is doable, and some folks at DeepMind may already do this. It might proceed slightly differently though, like

1. define a computation in jax
2. get the XLA HLO program as a protobuf
3. hand the protobuf to C++, compile and call it using the XLA C++ APIs.

To get the protobuf, one can use jax.xla_computation. This file shows one way to do it (in that case meant to be run at bazel build time, so that a C++ library can bazel-depend on the generated proto). I'm not sure if we have an examples of the C++ side, but it must be using the API in these XLA files.

cc @jlebar @tomhennigan @hawkinsp in case they can provide more specific tips

Yeah! What @mattjj said.

Write your cold code in C++ and your inner loops in Python. It is glorious.

Just a little more detail on the XLA C++ client side... JAX itself (indirectly) uses the PJRT client (main header file here). It's built on top of the LocalClient, which you can also use directly, but I think PJRT will be handier because it includes extra runtime APIs for e.g. copying host data to devices and back. Feel free to reach out if you want help with the specifics.

This is awesome - I will take a first pass at this over the break :)