Will Cromar

Add profiler traces to `PjRtComputationClient`

6

The existing profiler works out of the box because the traces work the same way. Tested manually with TensorBoard: 

TPU Pod support with PjRt

10

- Move TPU-specific logic from `pjrt.py` to `tpu.py`. All of this logic is broadly applicable to all TPU VMs and isn't strictly related to PjRt. - Update `configure_topology` to support...

[`PjRtCApiClient` in upstream TensorFlow](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/pjrt/pjrt_c_api_client.h) The new PjRt C API will be used to enable the new TPU runtime that we will support long-term. Note: The `PjRtCApiClient` currently requires a Google-internal...

This PR brings the size of the GPU image down to `6.24 GB` from `17.2 GB` - Use local `xla` directory instead of cloning from github. This allows building an...

This script does not work any more and is no longer necessary.

## 🚀 Feature - Provide a consistent API for external packages to register available [PJRT C API Plugins](https://github.com/openxla/xla/blob/main/xla/pjrt/c/docs/pjrt_integration_guide.md). - TPU and GPU should use the same plugin API. ## Motivation...

Update XLA base image from 1.1 to 1.2

2

Follow up on https://github.com/pytorch/xla/pull/6952 to make sure CI passes with new workflow. cc @JackCaoG

# Description The test name currently has the suffix `-1vm-1vm`. Simplify the test config since there's no 2VM or XRT case to worry about anymore :confetti_ball: # Tests Diff shows...

# Description Please include a summary of relevant context/issue and your changes. # Tests Please describe the tests that you ran on TPUs to verify changes. **Instruction and/or command lines...