Clean up rank so that it uses Spinv and revert change in init.py

pytorch · May 14, 2024 · fec9793 · fec9793
1 parent 4eab1c3
commit fec9793
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 6 deletions.
diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp
@@ -3437,10 +3437,11 @@ static void linalg_lstsq_out_info(
     } else {
         auto [U, S, Vh] = at::_linalg_svd(input, false, true, "gesvd");
         rank = at::zeros({1}, at::kLong);
-        rank[0] = (S > rcond).sum();
+
         auto S_pinv = S.reciprocal();
         auto s1 = at::narrow(S, /*dim=*/-1, /*start=*/0, /*length=*/1);  // singular values are sorted in descending order
         S_pinv.masked_fill_(S < rcond * s1, 0);
+        rank[0] = (S != 0).sum();
         auto uhOther = at::matmul(U.adjoint(), other);
         if(S_pinv.dim() != uhOther.dim()) {
           S_pinv = S_pinv.unsqueeze(-1);
@@ -3458,7 +3459,7 @@ static void linalg_lstsq_out_info(
     bool compute_residuals = true;
     if (driver == "gelss" || driver == "gelsd") {
       if (input.dim() == 2) {
-        compute_residuals = (rank.item().toDouble() == n);
+        compute_residuals = (rank.item().toInt() == n);
       } else {
         // it is not clear what to do if some matrices have rank < n in case of batched input
         // For now let's compute the residuals only if all matrices have rank equal to n

diff --git a/third_party/ideep b/third_party/ideep
diff --git a/third_party/kineto b/third_party/kineto
diff --git a/torch/linalg/__init__.py b/torch/linalg/__init__.py
@@ -1032,7 +1032,7 @@
   - `'gelsd'` (tridiagonal reduction and SVD)
   - But if you run into memory issues: `'gelss'` (full SVD).
 
-For CUDA inputs, two drivers are available: 'gels' and 'gelss'.
+For CUDA input, the only valid driver is `'gels'`, which assumes that :attr:`A` is full-rank.
 
 See also the `full description of these drivers`_
 
@@ -1080,7 +1080,7 @@
 
 Keyword args:
     driver (str, optional): name of the LAPACK/MAGMA method to be used.
-        If `None`, `'gelsy'` is used for CPU inputs, `'gels'` and `'gelss'` for CUDA inputs.
+        If `None`, `'gelsy'` is used for CPU inputs and `'gels'` for CUDA inputs.
         Default: `None`.
 
 Returns: