Add API to assemble CPU shards to a sharded tensor

[jonb377]

This PR reintroduces #5630, which was reverted in #5680 due to failing CI on master.

The following patch shows the difference between this and the original PR:

diff --git a/torch_xla/csrc/init_python_bindings.cpp b/torch_xla/csrc/init_python_bindings.cpp
index fe18f9508..421066ba7 100644
--- a/torch_xla/csrc/init_python_bindings.cpp
+++ b/torch_xla/csrc/init_python_bindings.cpp
@@ -1720,8 +1720,8 @@ void InitXlaModuleBindings(py::module m) {
               << " vs " << expected_shard_shape;
         }
 
-        auto data_handle = WrapXlaData(ShardingUtil::CreateShardedData(
-            shards, local_devices, sharding_spec));
+        auto data_handle = ShardingUtil::CreateShardedData(
+            shards, local_devices, sharding_spec);
         XLATensorPtr xla_tensor = XLATensor::Create(std::move(data_handle));
         xla_tensor->SetShardingSpec(*sharding_spec);
         auto tensor = bridge::AtenFromXlaTensor(std::move(xla_tensor));

[alanwaketan]

LGTM.

[jonb377]

Kokoro failure is due to a dependency issue:

ERROR: Could not find a version that satisfies the requirement tf-nightly (from versions: none)
ERROR: No matching distribution found for tf-nightly

I'll merge after TPU CI. Thanks Jiewen!

[jonb377]

Looking into the TPU CI failure, that's new since the rebase. Passes locally on v4, it may be that the test breaks with 8 devices.

[jonb377]

Surprisingly the test actually failed on the original PR, but TPU CI still passed: https://github.com/pytorch/xla/runs/17442958115