Fix for API match

fix for api match

Author: sshonTT

torch_xla/csrc/init_python_bindings.cpp

@@ -761,9 +761,12 @@ std::string GetTensorsHloGraph(const std::vector<at::Tensor>& tensors,
 }
 
 std::optional<xla::OpSharding> GetXLAOpSharding(const at::Tensor& input) {
-  XLATensorPtr xtensor = bridge::GetXlaTensor(input);
+  auto xtensor = bridge::GetXlaTensor(input);
+  if (!xtensor.ok()) {
+    return std::nullopt;
+  }
   XLATensor::ShardingSpecPtr sharding_spec =
-      xtensor ? xtensor->sharding_spec() : nullptr;
+      xtensor.value() ? xtensor.value()->sharding_spec() : nullptr;
   if (sharding_spec != nullptr) {
     return sharding_spec->sharding;
   }
@@ -3350,8 +3353,10 @@ void InitXlaModuleBindings(py::module m) {
                std::string key = item.first.cast<std::string>();
                options[key] = py::str(item.second).cast<std::string>();
              }
-             runtime::GetComputationClientOrDie()->SetCustomCompileOptions(
-                 options);
+             XLA_ASSIGN_OR_THROW(
+                 runtime::ComputationClient * absl_nonnull const client,
+                 runtime::GetComputationClient());
+             client->SetCustomCompileOptions(options);
            })
       .def(
           // from an XLA tensor to a PyCapsule.

torch_xla/csrc/runtime/pjrt_computation_client.cpp

@@ -140,7 +140,6 @@ absl::Status PjRtComputationClient::Initialize() {
   auto tracked_devices = GetLocalDevices();
   tracked_devices.emplace_back(spmd_device_str);
   operation_manager_ = std::move(OperationManager(std::move(tracked_devices)));
-
   return absl::OkStatus();
 }
 

torch_xla/csrc/xla_sharding_util.cpp

@@ -897,13 +897,10 @@ xla::Shape ShardingUtil::GetAdjustedGlobalShape(const at::Tensor& tensor,
                                                 bool minibatch) {
   xla::Shape global_shape = CreateComputationShapeFromTensor(tensor, nullptr);
   if (minibatch) {
-    int num_local_devices =
-        runtime::GetComputationClientOrDie()->GetLocalDevices().size();
-    int num_global_devices =
-        runtime::GetComputationClientOrDie()->GetAllDevices().size();
-    XLA_CHECK(tile_assignment.size() == num_global_devices)
-        << "Minibatch sharding only supports sharding along the batch "
-           "dimension";
+    XLA_ASSIGN_OR_THROW(runtime::ComputationClient * absl_nonnull const client,
+                        runtime::GetComputationClient());
+    int num_local_devices = client->GetLocalDevices().size();
+    int num_global_devices = client->GetAllDevices().size();
     int batch_dim_shape =
         tensor.sizes()[0] * num_global_devices / num_local_devices;
     global_shape.set_dimensions(0, batch_dim_shape);