Part 1: Disambiguate custom sharding op for DeviceData IR nodes

Author: rpsilva-aws

torch_xla/csrc/xla_sharding_util.cpp

@@ -772,15 +772,11 @@ void ShardingUtil::XlaMarkSharding(const at::Tensor& input,
                         xtensor->shape(), static_cast<XlaDeviceType>(
                                               xtensor->GetDevice().type())));
 
-  // For Non DeviceData IR values, we directly attach the sharding spec
-  // to the xtensor.
-  const DeviceData* device_data_node = nullptr;
+  // For IR values, we directly attach the sharding spec to the xtensor. In case
+  // it is a DeviceData IR node, we similarly add a custom sharding IR node.
   if (xtensor->CurrentIrValue()) {
-    device_data_node = DeviceData::Cast(xtensor->CurrentIrValue().node.get());
-    if (!device_data_node) {
-      tensor_methods::custom_sharding_(xtensor, new_sharding_spec);
-      return;
-    }
+    tensor_methods::custom_sharding_(xtensor, new_sharding_spec);
+    return;
   }
 
   // For data, we need to deal with the data transfers between
@@ -816,11 +812,12 @@ void ShardingUtil::XlaMarkSharding(const at::Tensor& input,
     // tensor from the physical device to CPU. In that case, the value
     // must be present on the backend device.
     XLA_CHECK((xtensor->CurrentDataHandle() &&
-               xtensor->CurrentDataHandle()->HasValue()) ||
-              device_data_node != nullptr)
+               xtensor->CurrentDataHandle()->HasValue()))
         << "Cannot shard tensor. Data does not present on any device.";
     std::vector<XLATensorPtr> xla_tensors{xtensor};
-    cpu_tensor = XLAGraphExecutor::Get()->GetTensors(&xla_tensors)[0];
+    auto tensors = XLAGraphExecutor::Get()->GetTensors(&xla_tensors);
+    XLA_CHECK_EQ(tensors.size(), 1);
+    cpu_tensor = tensors[0];
   }
   auto xla_data = CreateTensorsData(
       std::vector<at::Tensor>{cpu_tensor},