Implement forward AD for linalg.svd and improve svd_backward (#70253)

Summary:
Pull Request resolved: #70253

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.

cc jianyuh nikitaved pearu mruberry walterddr IvanYashchuk xwang233 Lezcano

Test Plan: Imported from OSS

Reviewed By: mikaylagawarecki

Differential Revision: D33751982

Pulled By: mruberry

fbshipit-source-id: c2a4a92a921a732357e99c01ccb563813b1af512

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)
@@ -2347,10 +2350,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);
@@ -2417,7 +2417,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);
@@ -2732,7 +2732,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_linalg.py

@@ -2817,6 +2817,27 @@ def test_svd(self, device, dtype):
                 S_s = torch.svd(A, compute_uv=False).S
                 self.assertEqual(S_s, S)
 
+    @skipCUDAIfNoMagmaAndNoCusolver
+    @skipCPUIfNoLapack
+    @dtypes(torch.complex128)
+    def test_invariance_error_spectral_decompositions(self, device, dtype):
+        make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=True)
+        A = make_arg((3, 3))
+        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
+            U, _, Vh = torch.linalg.svd(A, full_matrices=False)
+            (U + Vh).sum().backward()
+
+        A = make_arg((3, 3))
+        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
+            V = torch.linalg.eig(A).eigenvectors
+            V.sum().backward()
+
+        A = make_arg((3, 3))
+        A = A + A.mH
+        with self.assertRaisesRegex(RuntimeError, "ill-defined"):
+            Q = torch.linalg.eigh(A).eigenvectors
+            Q.sum().backward()
+
     @skipCUDAIfNoCusolver  # MAGMA backend doesn't work in this case
     @skipCUDAIfRocm
     @precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4})

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

@@ -1471,7 +1471,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)