xla
TPU Pod support with PjRt
#3813
Merged
Overview
Commits
19
Changes
View On GitHub

TPU Pod support with PjRt #3813

will-cromar merged 19 commits into master from pjrt-tpu-pod
will-cromar
will-cromar
will-cromar will-cromar requested a review from JackCaoG JackCaoG 3 years ago
will-cromar will-cromar requested a review from miladm miladm 3 years ago
will-cromar will-cromar requested a review from yeounoh yeounoh 3 years ago
will-cromar will-cromar marked this pull request as ready for review 3 years ago
ronghanghu
will-cromar
will-cromar
ronghanghu
ronghanghu
will-cromar
ronghanghu
JackCaoG
ronghanghu
will-cromar Configure TPU topology for pods
5fb8b20f
will-cromar Use metadata for network endpoints
5d1a67e7
will-cromar Add TPU tests
653d7575
will-cromar Fix computing num local processes.
fb5155f2
will-cromar Add unit tests for tpu.py
22706096
will-cromar Remove unused imports
0ef59d0f
will-cromar formatting
cf64ce33
will-cromar Don't use metadata for process bounds on v3
00c50cb1
will-cromar Fix real device indices for v3
c902cc6e
will-cromar Add TPU unit test to run_tests.sh
792e03c9
will-cromar Update docstrings
43c64630
will-cromar Fix signatures of `num_processes` and `num_local_processes`.
57b9b0b2
will-cromar Make MeshShape more pythonic
2ea3017e
will-cromar Clarify naming of `num_processes`
3ce02c7c
will-cromar Don't try to get TPU task id on CPU
5747bd1c
will-cromar Don't use value of `mock.patch.dict` because CI uses 3.7
215c9c5e
will-cromar will-cromar force pushed from 4b9e2e5b to 215c9c5e 3 years ago
will-cromar Remove real devices test
c392ce36
JackCaoG
JackCaoG commented on 2022-08-10
will-cromar Address review comments
695e5df4
will-cromar fix docstring
d2272cd0
JackCaoG
JackCaoG approved these changes on 2022-08-12
will-cromar
will-cromar will-cromar merged 63c67dfc into master 3 years ago

Login to write a write a comment.

Login via GitHub

Reviewers
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Terms Of Use Privacy Policy Refund Policy Impressum