Improve eig tests in preparation for new eig backends

- Verify validity of eigenvectors using the eigen decomposition identity
  for improved robustness, as eigenvectors are not unique.
- Increases test reliability across backends (cuSOLVER, MAGMA, CPU).
- Tolerances derived from numerical comparisons between cuSOLVER and NumPy. See discussion: https://dev-discuss.pytorch.org/t/cusolver-dnxgeev-faster-cuda-eigenvalue-calculations/3248/6

Author: johannesz-codes

test/test_linalg.py

@@ -2139,6 +2139,34 @@ def run_test(shape, *, symmetric=False):
     @skipCUDAIfNoMagma
     @dtypes(*floating_and_complex_types())
     def test_eig_compare_backends(self, device, dtype):
+
+        def fulfills_eigen_decomposition_identity(a, eig, dtype):
+
+            # check correctness using eigendecomposition identity
+            if dtype in [torch.float32, torch.complex64]:
+                atol = 1e-3  # CUDA seems to give less accurate results for float32 and complex64. They are about one to two OOM from NumPy results
+            else:
+                atol = 1e-13  # Same OOM for NumPy
+
+
+            w, v = eig
+
+            # move all tensors to CPU to avoid issues with GPU matmul
+            a = a.cpu().to(v.dtype)
+            v = v.to("cpu")
+            w = w.to("cpu")
+
+            if a.numel() == 0 and v.numel() == 0 and w.numel() == 0:
+                return True
+            elif a.numel() == 0 or v.numel() == 0 or w.numel() == 0:
+                return False
+
+            diff = (a @ v) - (v * w.unsqueeze(-2))
+            diff = diff.abs()
+            diff = torch.max(diff)
+            print(f"diff: {diff}")
+            return diff <= atol
+
         def run_test(shape, *, symmetric=False):
             from torch.testing._internal.common_utils import random_symmetric_matrix
 
@@ -2155,7 +2183,7 @@ def run_test(shape, *, symmetric=False):
             # compare with CPU
             expected = torch.linalg.eig(a.to(complementary_device))
             self.assertEqual(expected[0], actual[0])
-            self.assertEqual(expected[1], actual[1])
+            self.assertTrue(fulfills_eigen_decomposition_identity(a, actual, dtype))  # check evs using eigen identity
 
         shapes = [(0, 0),  # Empty matrix
                   (5, 5),  # Single matrix
@@ -2345,6 +2373,8 @@ def run_test(shape, *, symmetric=False):
             run_test(shape)
             run_test(shape, symmetric=True)
 
+
+
     @skipCUDAIfNoMagma
     @skipCPUIfNoLapack
     @dtypes(*floating_and_complex_types())