Experimental TPU Pod support with PjRt (#3813)
* Configure TPU topology for pods
* Use metadata for network endpoints
* Add TPU tests
* Fix computing num local processes.
* Add unit tests for tpu.py
* Add TPU unit test to run_tests.sh
* Update docstrings
