Kavish Gandhi

Create and Expose the `torch_xla::OpSharding` wrapper class instead of `xla::OpSharding` class

2

As a part of localized SPMD/submeshing effort we need abstract `xla::OpSharding` proto object with torch_xla specific wrapper class and expose the torch_xla::OpSharding object instead of xla proto object to the...

Changes needed for `device_assignment` in `PjRtComputationClient::Compile` to support submeshing/localized spmd

1

Currently `xla::DeviceAssignment` used inside `PjRtComputationClient::Compile` call makes use of the local devices/device_count for device assignment, but for submeshing or even for localized SPMD we would want to make use of...

Currently users cannot define the mesh with device_ids starting from anything except 0. This blocks us from defining sub-meshes, and also blocks the user from using localized SPMD within a...

feat: abstraction of xla::OpSharding proto using wrapper class

1

This PR includes the changes related to abstracting `xla::OpSharidng` proto object into a `torch_xla::OpSharding` wrapper class. This new class object will not have the requirements of xla::OpSharding (however, it will...

Adding functionality for metadata communication across hosts

1

This PR implement pipeline parallelism support for XLA devices with cross-host metadata communication capabilities. And has been tested on NEURON devices. Key points - - Feature addition implementing cross-host metadata...

This issue stems out from the [comment](https://github.com/pytorch/xla/pull/9467#discussion_r2286444496) on PR [feat: abstraction of xla::OpSharding proto using wrapper class](https://github.com/pytorch/xla/pull/9467#top) #9467. The aim is to move the constructor for the `IfrtComputation` (and similarly...