Implement forward AD for linalg.svd and improve svd_backward

I included a derivation of the formula in the complex case, as it is
particularly tricky. As far as I know, this is the first time this formula
is derived in the literature.

I also implemented a more efficient and more accurate version of svd_backward.
More importantly, I also added a lax check in the complex case making sure the loss
function just depends on the subspaces spanned by the pairs of singular
vectors, and not their joint phase.

ghstack-source-id: 2732ea3
Pull Request resolved: #70253

Author: lezcano

aten/src/ATen/native/LinearAlgebra.cpp

@@ -285,6 +285,9 @@ Tensor linalg_pinv(
     const optional<Tensor>& atol_opt,
     const optional<Tensor>& rtol_opt,
     bool hermitian) {
+  // FIXME: Whenever we have a nice lstsq, we should dispatch this function to simply be
+  // `torch.lstsq(A, torch.eye(A.shape[-1]), atol=atol, rtol=rtol)`
+  // with a driver that supports singular inputs
   NoTF32Guard disable_tf32;
   ScalarType t = input.scalar_type();
   TORCH_CHECK((t == ScalarType::Double || t == ScalarType::Float || t == ScalarType::ComplexFloat || t == ScalarType::ComplexDouble)
@@ -2342,10 +2345,7 @@ Tensor& nuclear_norm_out(const Tensor& self, IntArrayRef dim, bool keepdim, Tens
 
   auto permutation = create_dim_backshift_permutation(dim_[0], dim_[1], self.dim());
   Tensor p = self.permute(permutation);
-  // NOTE: U and V are computed only if gradmode is enabled, since the backward for nuclear
-  //       norm uses svd_backward, which requires them.
-  Tensor result_ = at::sum(std::get<1>(at::svd(p, /*some=*/true,
-                  /*compute_uv=*/at::GradMode::is_enabled() && self.requires_grad())), -1, keepdim);
+  Tensor result_ = at::sum(at::linalg_svdvals(p), -1, keepdim);
   if (keepdim) {
     result_.unsqueeze_(-1);
     auto permutation_reverse = create_reverse_permutation(permutation);
@@ -2412,7 +2412,7 @@ static Tensor& _linalg_norm_matrix_out(Tensor& result, const Tensor &self, const
   }
 
   if (std::abs(ord) == 2) {
-    // Need to shift the reduction dims to the back, because at::svd will only operate on
+    // Need to shift the reduction dims to the back, because at::linalg_svdvals will only operate on
     // the last 2 dimensions
     auto permutation = create_dim_backshift_permutation(dim_[0], dim_[1], self.dim());
     auto permutation_reverse = create_reverse_permutation(permutation);
@@ -2727,7 +2727,7 @@ Tensor linalg_cond(const Tensor& self, const optional<Scalar>& opt_ord) {
 
   // If ord == None or ord == ±2
   if (std::abs(ord.toDouble()) == 2.0) {
-    auto singular_values = std::get<1>(at::svd(self));
+    auto singular_values = at::linalg_svdvals(self);
     // singular values are sorted in descending order
     auto s_max = at::narrow(singular_values, /*dim=*/-1, /*start=*/0, /*length=*/1);
     auto s_min = at::narrow(singular_values, /*dim=*/-1, /*start=*/-1, /*length=*/1);

test/test_ops.py

@@ -266,6 +266,9 @@ def test_noncontiguous_samples(self, device, dtype, op):
             if not test_grad:
                 continue
 
+            expected = sample_input.output_process_fn_grad(expected)
+            actual = sample_input.output_process_fn_grad(actual)
+
             if isinstance(expected, torch.Tensor):
                 expected_backward_tensor = expected
                 actual_backward_tensor = actual

tools/autograd/derivatives.yaml

@@ -1467,7 +1467,15 @@
 
 # We never call _linalg_svd with compute_uv=False in an autograd context, so we don't even consider it here
 - name: _linalg_svd(Tensor A, bool full_matrices=False, bool compute_uv=True) -> (Tensor U, Tensor S, Tensor Vh)
-  A: svd_backward(grads, full_matrices, U, S, Vh)
+  A: "svd_backward(full_matrices && grad_U.defined() ? grad_U.narrow(-1, 0, S.size(-1)) : grad_U,
+                   grad_S,
+                   full_matrices && grad_Vh.defined() ? grad_Vh.narrow(-2, 0, S.size(-1)) : grad_Vh,
+                   full_matrices ? U.narrow(-1, 0, S.size(-1)) : U,
+                   S,
+                   full_matrices ? Vh.narrow(-2, 0, S.size(-1)) : Vh)"
+  U: std::get<0>(linalg_svd_jvp(A_t, U, S, Vh, full_matrices))
+  S: std::get<1>(linalg_svd_jvp(A_t, U, S, Vh, full_matrices))
+  Vh: std::get<2>(linalg_svd_jvp(A_t, U, S, Vh, full_matrices))
 
 - name: symeig(Tensor self, bool eigenvectors=False, bool upper=True) -> (Tensor eigenvalues, Tensor eigenvectors)
   self: linalg_eig_backward(grads[0], grads[1], eigenvalues, eigenvectors_return, /*is_hermitian=*/true, /*symeig_eigenvector=*/eigenvectors)