This blog post have not utilised any generative technology (LLMs, etc), if you see any mistakes or have any suggestion contact the author here.
What is OpenXLA?
The openxla machine learning ecosystem is a collection of tools, compilers and reusable components used in machine learning compilation (both inference and training) at both large and small scales. Jax1 differentiable array programming library is not a part of openxla but relies on it heavily for model compilation, runtime offloading and multi-device training and inference. Most of the compiler and runtime ecosystem is built around the MLIR compiler infrastructure.
MLIR
Jax and Pallas
Jaxlib and XLA Python
StableHLO
IFRT, TFRT, and PJRT
IFRT Design
PJRT Design
XLA Compiler
XLA Stream Executor
IREE and inference at all scales
Conculusion
Footnotes
Jax relies heavily on Openxla projects such as stablehlo, pjrt, ifrt, around the xla compiler to the point where I would consider it a part of the greater openxla ecosystem.↩︎