As of June 28th, 2023, PyTorch/XLA now pulls XLA from OpenXLA. OpenXLA is an open source machine learning compiler XLA for GPUs, CPUs, and ML accelerators.
Previous to OpenXLA, PyTorch/XLA pulled XLA directly from TensorFlow. With our XLA to OpenXLA migration, PyTorch/XLA now pulls XLA from OpenXLA.
For PJRT runtime users, there is no change with this migration. For XRT runtime users, there is a separate XRT branch of PyTorch/XLA since OpenXLA doesn't support XRT.
Below is a performance visual comparison of throughput for ResNet50 pre and post the migration on different TPU hardwares.
| resnet50-pjrt-v2-8 | resnet50-pjrt-v4-8 | resnet50-pjrt-v4-32 | |
|---|---|---|---|
| Pre Migration | 18.59 | 20.06 | 27.92 |
| Post Migration | 18.63 | 19.94 | 27.14 |