We plan on creating this prototype based on the existing mark_sharding. We need an API such that each host TransferShardsToDevice for the devices they are associated with. We can consider the existing pipeline for mark_sharding:

mark_sharding -> _xla_mark_sharding -> XlaMarkSharding -> CreateTensorsData -> CreateShardedData -> TransferShardsToDevice

In our prototype, we will do something like:

create_global_tensor_from_local_process_data -> _create_global_tensor_from_local_process_data -> XlaGlobalTensorFromLocalProcessData -> CreateGlobalTensorData -> CreateGlobalShardedData -> TransferShardsToDevice

I am still working on this prototype; I wanted to give a follow-up in the progress:

I have done an attempt of prototyping by just trying to create CreateGlobalTensorData -> CreateGlobalShardedData -> TransferShardsToDevice, and creating something like load_local_shards_. I however ran into a series of issues related object creation.

I have increased the scope of this prototype for the entire flow cited in #8842 (comment).

I have ran into a couple type coversion issues with xla::Shape, but I currently believe there is a path forward by leveraging xla::ShapeUtil::MakeShape. This will hopefully unblock us to run some tests.