Handle empty inputs in XLA _linalg_svd

test/cpp/test_aten_xla_tensor_2.cpp

@@ -2,6 +2,7 @@
 
 #include <iostream>
 #include <tuple>
+#include <vector>
 
 #include <torch/torch.h>
 
@@ -429,6 +430,30 @@ TEST_F(AtenXlaTensorTest, TestLinalgSVD) {
   ExpectCounterChanged("xla::_linalg_svd", cpp_test::GetIgnoredCounters());
 }
 
+TEST_F(AtenXlaTensorTest, TestLinalgSVDEmpty) {
+  std::vector<std::vector<int64_t>> empty_sizes = {
+      {0, 0}, {0, 3}, {3, 0}, {2, 0, 3}, {2, 3, 0}};
+  for (const auto& sizes : empty_sizes) {
+    for (bool full_matrices : {false, true}) {
+      for (bool compute_uv : {false, true}) {
+        torch::Tensor a =
+            torch::empty(sizes, torch::TensorOptions(torch::kFloat));
+        auto expected = torch::_linalg_svd(a, full_matrices, compute_uv);
+        ForEachDevice([&](const torch::Device& device) {
+          torch::Tensor xla_a = CopyToDevice(a, device);
+          auto actual = torch::_linalg_svd(xla_a, full_matrices, compute_uv);
+          ASSERT_EQ(std::get<0>(expected).sizes(),
+                    std::get<0>(actual).sizes());
+          ASSERT_EQ(std::get<1>(expected).sizes(),
+                    std::get<1>(actual).sizes());
+          ASSERT_EQ(std::get<2>(expected).sizes(),
+                    std::get<2>(actual).sizes());
+        });
+      }
+    }
+  }
+}
+
 TEST_F(AtenXlaTensorTest, TestLinalgVectorNorm) {
   torch::Tensor a = torch::rand({4, 3}, torch::TensorOptions(torch::kFloat));
   std::vector<float> ords = {0.0, 1.5, std::numeric_limits<float>::infinity(),

torch_xla/csrc/aten_xla_type.cpp

@@ -4175,6 +4175,32 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> XLANativeFunctions::_linalg_svd(
     std::optional<std::string_view> /* driver */) {
   // The optional driver string is only for CUDA with a cuSOLVER backend.
   TORCH_LAZY_FN_COUNTER_TIMED_TRACING("xla::");
+  if (self.numel() == 0) {
+    const auto m = self.size(-2);
+    const auto n = self.size(-1);
+    auto singular_values_sizes = self.sizes().vec();
+    const auto k = std::min(m, n);
+    singular_values_sizes.pop_back();
+    singular_values_sizes.back() = k;
+    auto s = at::zeros(singular_values_sizes, self.options());
+
+    if (!compute_uv) {
+      auto u = at::zeros({0}, self.options());
+      auto vh = at::zeros({0}, self.options());
+      return std::make_tuple(std::move(u), std::move(s), std::move(vh));
+    }
+
+    auto u_sizes = self.sizes().vec();
+    u_sizes.back() = full_matrices ? m : k;
+    auto u = at::zeros(u_sizes, self.options());
+
+    auto vh_sizes = self.sizes().vec();
+    vh_sizes.end()[-2] = full_matrices ? n : k;
+    vh_sizes.end()[-1] = n;
+    auto vh = at::zeros(vh_sizes, self.options());
+
+    return std::make_tuple(std::move(u), std::move(s), std::move(vh));
+  }
   // As per https://pytorch.org/docs/stable/generated/torch.svd.html,
   // The second boolean argument is exactly opposite between
   // torch::svd and torch::_linalg_svd, hence the negation of full_matrices.