Prototype API to build global tensor using local processes

[pgmoka]

As a sub-item for enabling 2D sharding with minibatch=True, we want to create a protoype of SPMD that uses local batching. This will hopefully be the beginning of an implementation that might substitute the current workaround.

[pgmoka]

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

[pgmoka]

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.

[pgmoka]

I got pulled away last two weks from this bug. Current progress is this: the prototype seems to work, but I still need to validate it in a multi-host set-up.

After validating, we would need to clean-up the prototype, and create tests