fix for failing ci/cd tests

Author: kvshbg-aws

torch_xla/csrc/init_python_bindings.cpp

@@ -2632,12 +2632,12 @@ void InitXlaModuleBindings(py::module m) {
             return GetXLAShardingSpec(xtensor);
            })
       .def("_get_xla_op_sharding",
-           [](const at::Tensor& input) -> std::optional<xla::OpSharding> {
+           [](const at::Tensor& input) -> std::optional<torch_xla::OpSharding> {
             XLATensorPtr xtensor = GetValueOrThrow(bridge::GetXlaTensor(input));
             XLATensor::ShardingSpecPtr sharding_spec =
                 xtensor ? xtensor->sharding_spec() : nullptr;
             if (sharding_spec != nullptr) {
-              return sharding_spec->sharding.GetXlaOpSharding();
+              return sharding_spec->sharding;
             }
             return std::nullopt;
            })

torch_xla/csrc/runtime/ifrt_computation_client.h

@@ -287,7 +287,8 @@ class IfrtComputationClient : public ComputationClient {
               denormalized_tile_assignment.value_or(std::vector<int64_t>{}))) {
       xla_output_shardings_ = this->executable->GetOutputShardings();
       if (xla_output_shardings_.has_value()) {
-        output_shardings_->reserve(xla_output_shardings_->size());
+        output_shardings_ = std::vector<torch_xla::OpSharding>{};
+        output_shardings_->reserve(xla_output_shardings_.value().size());
         for (const auto& sharding : xla_output_shardings_.value()) {
           // convert each into torch_xla::OpSharding object
           torch_xla::OpSharding torch_xla_op_sharding(

torch_xla/csrc/runtime/pjrt_computation_client.h

@@ -347,7 +347,8 @@ class PjRtComputationClient : public ComputationClient {
               denormalized_tile_assignment.value_or(std::vector<int64_t>{}))) {
       xla_output_shardings_ = this->executable->GetOutputShardings();
       if (xla_output_shardings_.has_value()) {
-        output_shardings_->reserve(xla_output_shardings_->size());
+        output_shardings_ = std::vector<torch_xla::OpSharding>{};
+        output_shardings_->reserve(xla_output_shardings_.value().size());
         for (const auto& sharding : xla_output_shardings_.value()) {
           // convert each into torch_xla::OpSharding object
           torch_xla::OpSharding torch_xla_op_sharding(