PR #3813 TPU Pod support with PjRt

Configure TPU topology for pods

will-cromar committed 4 years ago

Use metadata for network endpoints

will-cromar committed 4 years ago

Add TPU tests

will-cromar committed 4 years ago

Fix computing num local processes.

will-cromar committed 4 years ago

Add unit tests for tpu.py

will-cromar committed 4 years ago

Remove unused imports

will-cromar committed 4 years ago

formatting

will-cromar committed 4 years ago

Don't use metadata for process bounds on v3

will-cromar committed 4 years ago

Fix real device indices for v3

will-cromar committed 4 years ago

Add TPU unit test to run_tests.sh

will-cromar committed 4 years ago

Update docstrings

will-cromar committed 4 years ago

Fix signatures of `num_processes` and `num_local_processes`.

will-cromar committed 4 years ago

Make MeshShape more pythonic

will-cromar committed 4 years ago

Clarify naming of `num_processes`

will-cromar committed 4 years ago

Don't try to get TPU task id on CPU

will-cromar committed 4 years ago

Don't use value of `mock.patch.dict` because CI uses 3.7

will-cromar committed 4 years ago

Remove real devices test

will-cromar committed 4 years ago

Address review comments

will-cromar committed 4 years ago

fix docstring

will-cromar committed 4 years ago

xla TPU Pod support with PjRt #3813 Merged