PR #3704 Add `torch_xla.experimental.pjrt`

Add `torch_xla.experimental.pjrt` #3704

will-cromar merged 17 commits into master from pjrt-xla-model

will-cromar force pushed from 3ce2f13e to e30ad73e 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

yeounoh requested a review from

yeounoh 3 years ago

JackCaoG commented on 2022-07-19

JackCaoG approved these changes on 2022-07-19

yeounoh commented on 2022-07-21

yeounoh approved these changes on 2022-07-21

Update get_ordinal, get_local_ordinal, and xla_device to work with PJRT

9dbb7326

Move common PJRT code to torch_xla.experimental.pjrt

91e20a2a

Add docstrings.

84fbc915

Add topology vars to xenv

b96799c7

Unit tests and type hints.

348d3be2

Better type hints

f7b042e8

formatting

a3345c0d

Add to CI tests.

f045a807

Set default device correctly

d5382ed2

Make test_set_ordinals clearer

93ed2570

Fix local ordinal and update naming

ecbd9680

Add test_world_size

86c329fe

Add _xla_num_devices to _XLAC

457d3485

Clarify docstrings

fe07ff95

Formatting

03414fb1

process -> thread

c286612b

will-cromar force pushed from abc22206 to c286612b 3 years ago

Address review comments.

0467d87a

will-cromar merged ad1fe4e3 into master 3 years ago

will-cromar deleted the pjrt-xla-model branch 3 years ago

Reviewers

JackCaoG

yeounoh

Assignees

No one assigned

Labels

None yet

Milestone

No milestone

xla Add `torch_xla.experimental.pjrt` #3704 Merged

Add `torch_xla.experimental.pjrt` #3704

xla
Add `torch_xla.experimental.pjrt`
#3704

Merged