xla
ad1fe4e3 - Add `torch_xla.experimental.pjrt` (#3704)

Commit
3 years ago
Add `torch_xla.experimental.pjrt` (#3704) * Update get_ordinal, get_local_ordinal, and xla_device to work with PJRT * Move common PJRT code to torch_xla.experimental.pjrt * Add topology vars to xenv * Unit tests and type hints. * Add to CI tests. * Set default device correctly * Add test_world_size * Add _xla_num_devices to _XLAC * Clarify docstrings
Author
Parents
Loading