Add test for hegvx

Author: Cstandardlib

source/source_base/module_container/ATen/kernels/test/lapack_test.cpp

@@ -188,6 +188,19 @@ TYPED_TEST(LapackTest, heevx) {
     // check that A*V = E*V
     E = E.to_device<DEVICE_CPU>();
     V = V.to_device<DEVICE_CPU>();
+    std::cout << "Eigenvalues E: ";
+    for (int i = 0; i < neig; i++) {
+        std::cout << E.data<Real>()[i] << " ";
+    }
+    std::cout << std::endl;
+
+    std::cout << "Eigenvectors V:" << std::endl;
+    for (int i = 0; i < dim; i++) {
+        for (int j = 0; j < neig; j++) {
+            std::cout << V.data<Type>()[i + j * dim] << " ";
+        }
+        std::cout << std::endl;
+    }
 
     EXPECT_EQ(expected_C1, expected_C2);
 }
@@ -242,7 +255,69 @@ TYPED_TEST(LapackTest, hegvd) {
     EXPECT_EQ(expected_C1, expected_C2);
 }
 
+TYPED_TEST(LapackTest, hegvx) {
+    using Type = typename std::tuple_element<0, decltype(TypeParam())>::type;
+    using Real = typename GetTypeReal<Type>::type;
+    using Device = typename std::tuple_element<1, decltype(TypeParam())>::type;
 
+    blas_gemm<Type, Device> gemmCalculator;
+    blas_axpy<Type, Device> axpyCalculator;
+    lapack_hegvx<Type, Device> hegvxCalculator;
+
+    const int dim = 3;
+    const int neig = 2;  // Compute first 2 eigenvalues
+
+    Tensor A = std::move(Tensor({static_cast<Type>(4.0), static_cast<Type>(1.0), static_cast<Type>(1.0),
+                                 static_cast<Type>(1.0), static_cast<Type>(5.0), static_cast<Type>(3.0),
+                                 static_cast<Type>(1.0), static_cast<Type>(3.0), static_cast<Type>(6.0)}).to_device<Device>());
+
+    Tensor B = std::move(Tensor({static_cast<Type>(2.0), static_cast<Type>(0.0), static_cast<Type>(0.0),
+                                 static_cast<Type>(0.0), static_cast<Type>(2.0), static_cast<Type>(0.0),
+                                 static_cast<Type>(0.0), static_cast<Type>(0.0), static_cast<Type>(2.0)}).to_device<Device>());
+
+    Tensor E = std::move(Tensor({static_cast<Real>(0.0), static_cast<Real>(0.0)}).to_device<Device>());
+    Tensor V = A;
+    Tensor expected_C1 = std::move(Tensor({static_cast<Type>(0.0), static_cast<Type>(0.0), static_cast<Type>(0.0),
+                                           static_cast<Type>(0.0), static_cast<Type>(0.0), static_cast<Type>(0.0)}).to_device<Device>());
+    Tensor expected_C2 = expected_C1;
+    Tensor C_temp = expected_C1;
+    expected_C1.zero();
+    expected_C2.zero();
+
+    const char trans = 'N';
+    const int m = 3;
+    const int n = neig;
+    const int k = 3;
+    const Type alpha = static_cast<Type>(1.0);
+    const Type beta  = static_cast<Type>(0.0);
+
+    // Compute first neig eigenvalues and eigenvectors using hegvx
+    hegvxCalculator(dim, dim, A.data<Type>(), B.data<Type>(), neig, E.data<Real>(), V.data<Type>());
+
+    E = E.to_device<ct::DEVICE_CPU>();
+    const Tensor Alpha = std::move(Tensor({
+            static_cast<Type>(E.data<Real>()[0]),
+            static_cast<Type>(E.data<Real>()[1])}));
+
+    // Check the eigenvalues and eigenvectors
+    // A * x = lambda * B * x for the first neig eigenvectors
+    // get A*V
+    gemmCalculator(trans, trans, m, n, k, &alpha, A.data<Type>(), m, V.data<Type>(), k, &beta, expected_C1.data<Type>(), m);
+    // get E * B * V
+    // where B is 2 * eye(3,3)
+    // get C_temp = B * V first
+    gemmCalculator(trans, trans, m, n, k, &alpha, B.data<Type>(), m, V.data<Type>(), k, &beta, C_temp.data<Type>(), m);
+    // then compute C2 = E * B * V
+    for (int ii = 0; ii < neig; ii++) {
+        axpyCalculator(dim, Alpha.data<Type>() + ii, C_temp.data<Type>() + ii * dim, 1, expected_C2.data<Type>() + ii * dim, 1);
+    }
+    // check that A*V = E*V
+    E = E.to_device<DEVICE_CPU>();
+    V = V.to_device<DEVICE_CPU>();
+
+
+    EXPECT_EQ(expected_C1, expected_C2);
+}
 
 } // namespace kernels
 } // namespace container