Accelerator Software Stack

XLA / JAX

google

Accelerated Linear Algebra compiler used by JAX and TensorFlow for TPUs

Editorial compatibility: these components can be used together in representative deployments. That's not a requirement or a validated implementation. For real deployments, see the sourced reference architectures.

Compatible from earlier layers

Ray Serve

Compatible with later layers

TPU v5e
TPU v6e (Trillium)