xla
Experimental support for single-host TPU in PjRt
#3550

Merged

Experimental support for single-host TPU in PjRt #3550

JackCaoG merged 8 commits into master from wcromar-pjrt-tpu

will-cromar added runtime

will-cromar force pushed from 0be7da38 to 714ced10 3 years ago

yeounoh requested a review from

yeounoh 3 years ago

will-cromar requested a review from

JackCaoG 3 years ago

will-cromar marked this pull request as ready for review 3 years ago

JackCaoG commented on 2022-05-10

Experimental support for single-device TPU in PjRt.

806a2dee

Distribute TPU execution across multiple replicas.

4e201766

Fix compilation device.

05a57edf

Formatting

bbd53f41

string_to_device -> string_to_device_

829493ce

Fix issues found in model testing

bc44d516

Disable strict shape checking.

62c4c5a4

will-cromar force pushed from 2011a507 to 62c4c5a4 3 years ago

Use absl::AsciiStrToUpper instead of std::transform

04d637a0

will-cromar requested a review from

JackCaoG 3 years ago

JackCaoG approved these changes on 2022-05-13

JackCaoG merged 1ab84187 into master 3 years ago

JackCaoG deleted the wcromar-pjrt-tpu branch 3 years ago

yeounoh commented on 2022-05-17

Reviewers

JackCaoG

yeounoh

Assignees

No one assigned

Labels

runtime

Milestone

No milestone