Update test_eigenvalues in test_linalg.py

Author: vlad-perevezentsev

dpnp/linalg/dpnp_utils_linalg.py

@@ -932,6 +932,14 @@ def dpnp_eigh(a, UPLO, eigen_mode="V"):
                 sycl_queue=a_sycl_queue,
             )
 
+            # TODO: Remove this w/a when MKLD-17201 is solved.
+            # Waiting for a host task executing an OneMKL LAPACK syevd call
+            # on CPU causes deadlock due to serialization of all host tasks
+            # in the queue.
+            # We need to wait for each host tasks before calling _seyvd to avoid deadlock.
+            if lapack_func == "_syevd" and is_cpu_device:
+                ht_list_ev[2 * i].wait()
+
             # call LAPACK extension function to get eigenvalues and eigenvectors of a portion of matrix A
             ht_list_ev[2 * i + 1], _ = getattr(li, lapack_func)(
                 a_sycl_queue,
@@ -942,15 +950,6 @@ def dpnp_eigh(a, UPLO, eigen_mode="V"):
                 depends=[copy_ev],
             )
 
-            # TODO: Remove this w/a when MKLD-17201 is solved.
-            # Waiting for a host task executing an OneMKL LAPACK syevd call
-            # on CPU causes deadlock due to serialization of all host tasks
-            # in the queue.
-            # We need to wait for each host tasks before calling _seyvd again
-            # to avoid deadlock.
-            if lapack_func == "_syevd" and is_cpu_device:
-                ht_list_ev[2 * i + 1].wait()
-
         dpctl.SyclEvent.wait_for(ht_list_ev)
 
         w = w.reshape(w_orig_shape)

tests/test_linalg.py

@@ -386,6 +386,23 @@ def test_eig_arange(type, size):
 
 
 class TestEigenvalue:
+    # Eigenvalue decomposition of a matrix or a batch of matrices
+    # by checking if the eigen equation A*v=w*v holds for given eigenvalues(w)
+    # and eigenvectors(v).
+    def assert_eigen_decomposition(self, a, w, v, rtol=1e-5, atol=1e-5):
+        a_ndim = a.ndim
+        if a_ndim == 2:
+            assert_allclose(a @ v, v @ inp.diag(w), rtol=rtol, atol=atol)
+        else:  # a_ndim > 2
+            if a_ndim > 3:
+                a = a.reshape(-1, *a.shape[-2:])
+                w = w.reshape(-1, w.shape[-1])
+                v = v.reshape(-1, *v.shape[-2:])
+            for i in range(a.shape[0]):
+                assert_allclose(
+                    a[i].dot(v[i]), w[i] * v[i], rtol=rtol, atol=atol
+                )
+
     @pytest.mark.parametrize(
         "func",
         [
@@ -413,11 +430,16 @@ def test_eigenvalues(self, func, shape, dtype, order):
         a_order = numpy.array(a, order=order)
         a_dp = inp.array(a, order=order)
 
+        # NumPy with OneMKL and with rocSOLVER sorts in ascending order,
+        # so w's should be directly comparable.
+        # However, both OneMKL and rocSOLVER pick a different convention for
+        # constructing eigenvectors, so v's are not directly comparible and
+        # we verify them through the eigen equation A*v=w*v.
         if func == "eigh":
-            w, v = numpy.linalg.eigh(a_order)
+            w, _ = numpy.linalg.eigh(a_order)
             w_dp, v_dp = inp.linalg.eigh(a_dp)
 
-            assert_dtype_allclose(v_dp, v)
+            self.assert_eigen_decomposition(a_dp, w_dp, v_dp)
 
         else:  # eighvalsh
             w = numpy.linalg.eigvalsh(a_order)