Coral TPU

[soldierofhell]

Hi, can we deploy JAX "models" on Coral TPU like
https://coral.ai/products/accelerator/ ? Regards

[skye]

It looks like the Coral TPU only interfaces with Tensorflow Lite (https://coral.ai/docs/edgetpu/inference/). It may be possible to use jax2tf to convert a JAX model to a SavedModel, which I'm guessing can be fed into TFLite. @gnecula may know if anyone's tried this before.

[gnecula]

I am not aware of anybody having tried that, but it is worth a try. Essentially jax2tf adds a new box at the start of the deployment diagram to create a TF model from a JAX one. Note that you will have to use enable_xla=False in jax2tf.convert to make it more likely that the TFlite converter will like your model. There is an example available.

[soldierofhell]

Thanks. I was hoping that there's some "native" JAX option (with all jax "ops" supported). Do you know if this TFLite-only is hardware-limited or it's possible to run JAX but just there's no implemented path to do it. In other words is this Coral TPU a limited-to-TFLite version of "real" TPU?

[skye]

I'm not sure about any of this, but I think it would be possible to target Coral TPU with JAX with only software changes. However, it'd be a pretty big software change. My understanding is that TFLite takes a TF graph and generates machine code for mobile or edge devices, including the Coral TPU. JAX uses the XLA compiler to generate machine code for CPU, GPU, or (non-edge) TPU. I don't think XLA knows how to target the Coral TPU, but it could theoretically be extended to do so. That's how you would get "native" JAX support.

Maybe reach out to the Coral TPU people too, they probably know more!