Running JAX code on Graphcore IPUs

[jacobthebanana]

Hello, is it possible to run JAX on Graphcore's IPUs? Is there a particular jaxlib version that allows JAX to run on Graphcore's IPUs? If not, how much work would it take to build an IPU version of jaxlib? Graphcore IPUs do provide an XLA backend with proper support for TensorFlow, though it's unclear to me whether JAX could make use of that backend.

The goal here is to take advantage of the free IPU machines (4 petaFLOPS FP16 each) on Paperspace: paperspace.com/graphcore without having to write any Graphcore-specific code.