Integrating JAX with a custom accelerator

[slai-nick]

Is there any path for integrating JAX with a custom accelerator?
Maybe by providing a custom JIT compiler? Although I don't think that's enough to handle cases of memory transfer between devices.

If not, is it on the roadmap to add such a feature?

[hawkinsp]

I'm assuming here you have say, a novel GPU or TPU-like accelerator device that you would like to integrate into JAX. Yes, this is something that is on the roadmap.

The exact mechanics of how you do this will change, but in essence one implements new JAX devices by providing a compiler and runtime that implement this API:

https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/compiler/xla/pjrt/pjrt_client.h

The PJRT client API is what JAX uses to compile and run MHLO programs and to perform other runtime tasks, such as management of buffers.

Currently such an extension also requires building and shipping a custom jaxlib into which your custom plugin is integrated, however we are working on making that API into a more pluggable C API.

[slai-nick]

Thank you for the clear answer!