Added linalg.eigh, linalg.eigvalsh

aten/src/ATen/native/BatchLinearAlgebra.cpp

@@ -71,6 +71,12 @@ extern "C" void cheev_(char *jobz, char *uplo, int *n, std::complex<float> *a, i
 extern "C" void dsyev_(char *jobz, char *uplo, int *n, double *a, int *lda, double *w, double *work, int *lwork, int *info);
 extern "C" void ssyev_(char *jobz, char *uplo, int *n, float *a, int *lda, float *w, float *work, int *lwork, int *info);
 
+// syevd
+extern "C" void zheevd_(char *jobz, char *uplo, int *n, std::complex<double> *a, int *lda, double *w, std::complex<double> *work, int *lwork, double *rwork, int *lrwork, int *iwork, int *liwork, int *info);
+extern "C" void cheevd_(char *jobz, char *uplo, int *n, std::complex<float> *a, int *lda, float *w, std::complex<float> *work, int *lwork, float *rwork, int *lrwork, int *iwork, int *liwork, int *info);
+extern "C" void dsyevd_(char *jobz, char *uplo, int *n, double *a, int *lda, double *w, double *work, int *lwork, int *iwork, int *liwork, int *info);
+extern "C" void ssyevd_(char *jobz, char *uplo, int *n, float *a, int *lda, float *w, float *work, int *lwork, int *iwork, int *liwork, int *info);
+
 // gesdd
 extern "C" void zgesdd_(char *jobz, int *m, int *n, std::complex<double> *a, int *lda,
                         double *s, std::complex<double> *u, int *ldu, std::complex<double> *vt, int *ldvt, std::complex<double> *work, int *lwork, double *rwork, int *iwork, int *info);
@@ -121,6 +127,9 @@ void lapackOrgqr(int m, int n, int k, scalar_t *a, int lda, scalar_t *tau, scala
 template<class scalar_t, class value_t=scalar_t>
 void lapackSymeig(char jobz, char uplo, int n, scalar_t *a, int lda, value_t *w, scalar_t *work, int lwork, value_t *rwork, int *info);
 
+template<class scalar_t, class value_t=scalar_t>
+void lapackSyevd(char jobz, char uplo, int n, scalar_t *a, int lda, value_t *w, scalar_t *work, int lwork, value_t *rwork, int lrwork, int *iwork, int liwork, int *info);
+
 template<class scalar_t, class value_t=scalar_t>
 void lapackSvd(char jobz, int m, int n, scalar_t *a, int lda,
                value_t *s, scalar_t *u, int ldu, scalar_t *vt, int ldvt, scalar_t *work, int lwork, value_t *rwork, int *iwork, int *info);
@@ -275,6 +284,26 @@ template<> void lapackSymeig<float>(char jobz, char uplo, int n, float *a, int l
   ssyev_(&jobz, &uplo, &n, a, &lda, w, work, &lwork, info);
 }
 
+template<> void lapackSyevd<c10::complex<double>, double>(char jobz, char uplo, int n, c10::complex<double> *a, int lda, double *w, c10::complex<double> *work, int lwork, double *rwork, int lrwork, int *iwork, int liwork, int *info) {
+  zheevd_(&jobz, &uplo, &n, reinterpret_cast<std::complex<double>*>(a), &lda, w, reinterpret_cast<std::complex<double>*>(work), &lwork, rwork, &lrwork, iwork, &liwork, info);
+}
+
+template<> void lapackSyevd<c10::complex<float>, float>(char jobz, char uplo, int n, c10::complex<float> *a, int lda, float *w, c10::complex<float> *work, int lwork, float *rwork, int lrwork, int *iwork, int liwork, int *info) {
+  cheevd_(&jobz, &uplo, &n, reinterpret_cast<std::complex<float>*>(a), &lda, w, reinterpret_cast<std::complex<float>*>(work), &lwork, rwork, &lrwork, iwork, &liwork, info);
+}
+
+template<> void lapackSyevd<double>(char jobz, char uplo, int n, double *a, int lda, double *w, double *work, int lwork, double *rwork, int lrwork, int *iwork, int liwork, int *info) {
+  (void)rwork;  // unused
+  (void)lrwork;  // unused
+  dsyevd_(&jobz, &uplo, &n, a, &lda, w, work, &lwork, iwork, &liwork, info);
+}
+
+template<> void lapackSyevd<float>(char jobz, char uplo, int n, float *a, int lda, float *w, float *work, int lwork, float *rwork, int lrwork, int *iwork, int liwork, int *info) {
+  (void)rwork;  // unused
+  (void)lrwork;  // unused
+  ssyevd_(&jobz, &uplo, &n, a, &lda, w, work, &lwork, iwork, &liwork, info);
+}
+
 template<> void lapackSvd<c10::complex<double>, double>(char jobz, int m, int n, c10::complex<double> *a, int lda,
                                   double *s, c10::complex<double> *u, int ldu, c10::complex<double> *vt, int ldvt, c10::complex<double> *work, int lwork, double *rwork, int *iwork, int *info) {
   zgesdd_(&jobz, &m, &n, reinterpret_cast<std::complex<double>*>(a), &lda, s, reinterpret_cast<std::complex<double>*>(u), &ldu,
@@ -862,6 +891,101 @@ std::tuple<Tensor&,Tensor&> qr_out(Tensor& Q, Tensor& R, const Tensor& self, boo
   return std::tuple<Tensor&, Tensor&>(Q, R);
 }
 
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ syevd ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+template <typename scalar_t>
+static void apply_syevd(Tensor& w, Tensor& v, bool compute_v, std::string uplo_str, std::vector<int64_t>& infos) {
+#ifndef USE_LAPACK
+  AT_ERROR("syevd: LAPACK library not found in compilation");
+#else
+  using value_t = typename c10::scalar_value_type<scalar_t>::type;
+
+  auto v_data = v.data_ptr<scalar_t>();
+  auto w_data = w.data_ptr<value_t>();
+  auto v_matrix_stride = matrixStride(v);
+  auto w_stride = w.size(-1);
+  auto batch_size = batchCount(v);
+  auto n = v.size(-1);
+  auto lda = std::max(int64_t{1}, n);
+
+  char uplo = uplo_str == "U" ? 'U' : 'L';
+  char jobz = compute_v ? 'V' : 'N';
+
+  int info;
+  // Run once, first to get the optimum work size.
+  // Since we deal with batches of matrices with the same dimensions, doing this outside
+  // the loop saves (batch_size - 1) workspace queries which would provide the same result
+  // and (batch_size - 1) calls to allocate and deallocate workspace using at::empty()
+  int lwork = -1;
+  int lrwork = -1;
+  int liwork = -1;
+  scalar_t work_query;
+  value_t rwork_query;
+  int iwork_query;
+
+  lapackSyevd<scalar_t, value_t>(jobz, uplo, n, v_data, lda, w_data, &work_query, lwork, &rwork_query, lrwork, &iwork_query, liwork, &info);
+  lwork = std::max(1, static_cast<int>(real_impl<scalar_t, value_t>(work_query)));
+  Tensor work = at::empty({lwork}, v.options());
+  liwork = std::max(1, static_cast<int>(iwork_query));
+  Tensor iwork = at::empty({liwork}, at::kInt);
+
+  Tensor rwork;
+  value_t* rwork_data = nullptr;
+  if (isComplexType(at::typeMetaToScalarType(v.dtype()))) {
+    lrwork = std::max(1, static_cast<int>(rwork_query));
+    rwork = at::empty({lrwork}, w.options());
+    rwork_data = rwork.data_ptr<value_t>();
+  }
+
+  for (int64_t i = 0; i < batch_size; i++) {
+    scalar_t* v_working_ptr = &v_data[i * v_matrix_stride];
+    value_t* w_working_ptr = &w_data[i * w_stride];
+    lapackSyevd<scalar_t, value_t>(jobz, uplo, n, v_working_ptr, lda, w_working_ptr, work.data_ptr<scalar_t>(), lwork, rwork_data, lrwork, iwork.data_ptr<int>(), liwork, &info);
+    infos[i] = info;
+    if (info != 0) {
+      return;
+    }
+  }
+#endif
+}
+
+std::tuple<Tensor, Tensor> _syevd_helper_cpu(const Tensor& self, bool compute_v, std::string uplo) {
+  std::vector<int64_t> infos(batchCount(self), 0);
+
+  auto self_sizes = self.sizes().vec();
+  self_sizes.pop_back();
+  ScalarType dtype = toValueType(typeMetaToScalarType(self.dtype()));
+  auto eigvals = at::empty(self_sizes, self.options().dtype(dtype));
+
+  auto eigvecs = cloneBatchedColumnMajor(self);
+  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES(self.scalar_type(), "syevd_cpu", [&]{
+    apply_syevd<scalar_t>(eigvals, eigvecs, compute_v, uplo, infos);
+  });
+
+  if (self.dim() > 2) {
+    batchCheckErrors(infos, "syevd_cpu");
+  } else {
+    singleCheckErrors(infos[0], "syevd_cpu");
+  }
+  if (compute_v) {
+    return std::tuple<Tensor, Tensor>(eigvals, eigvecs);
+  } else {
+    return std::tuple<Tensor, Tensor>(eigvals, at::empty({0}, self.options()));
+  }
+}
+
+std::tuple<Tensor, Tensor> linalg_eigh(const Tensor& self, std::string uplo) {
+  squareCheckInputs(self);
+  return at::_syevd_helper(self, /*compute_v=*/true, uplo);
+}
+
+Tensor linalg_eigvalsh(const Tensor& self, std::string uplo) {
+  squareCheckInputs(self);
+  Tensor eigvals, eigvecs;
+  std::tie(eigvals, eigvecs) = at::_syevd_helper(self, /*compute_v=*/false, uplo);
+  return eigvals;
+}
+
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ symeig ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 template <typename scalar_t>

aten/src/ATen/native/cuda/BatchLinearAlgebra.cu

@@ -1389,6 +1389,13 @@ std::tuple<Tensor, Tensor> _symeig_helper_cuda(const Tensor& self, bool eigenvec
   }
 }
 
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ syevd ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+std::tuple<Tensor, Tensor> _syevd_helper_cuda(const Tensor& self, bool compute_v, std::string uplo) {
+  bool upper = uplo == "U" ? true : false;
+  return _symeig_helper_cuda(self, compute_v, upper);
+}
+
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ svd ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 template<typename scalar_t>

aten/src/ATen/native/native_functions.yaml

@@ -8224,6 +8224,23 @@
   use_c10_dispatcher: full
   variants: function, method
 
+- func: _syevd_helper(Tensor self, bool compute_v, str uplo) -> (Tensor, Tensor)
+  use_c10_dispatcher: full
+  variants: function
+  dispatch:
+    CPU: _syevd_helper_cpu
+    CUDA: _syevd_helper_cuda
+
+- func: linalg_eigh(Tensor self, str uplo) -> (Tensor, Tensor)
+  python_module: linalg
+  use_c10_dispatcher: full
+  variants: function
+
+- func: linalg_eigvalsh(Tensor self, str uplo) -> Tensor
+  python_module: linalg
+  use_c10_dispatcher: full
+  variants: function
+
 # torch.outer, alias for torch.ger
 - func: outer(Tensor self, Tensor vec2) -> Tensor
   use_c10_dispatcher: full

docs/source/linalg.rst

@@ -13,4 +13,6 @@ Functions
 ---------
 
 .. autofunction:: det
+.. autofunction:: eigh
+.. autofunction:: eigvalsh
 .. autofunction:: norm

test/test_linalg.py

@@ -180,6 +180,54 @@ def test_det(self, device, dtype):
         with self.assertRaises(RuntimeError):
             op(t)
 
+    @skipCUDAIfNoMagma
+    @skipCPUIfNoLapack
+    @unittest.skipIf(not TEST_NUMPY, "NumPy not found")
+    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
+    def test_eigh(self, device, dtype):
+        from torch.testing._internal.common_utils import random_hermitian_matrix
+
+        def run_test(shape, batch):
+            matrix = random_hermitian_matrix(shape, *batch, dtype=dtype, device=device)
+            expected_w, expected_v = np.linalg.eigh(matrix.cpu().numpy())
+            actual_w, actual_v = torch.linalg.eigh(matrix)
+            self.assertEqual(actual_w, expected_w)
+            # sign of eigenvectors is not unique and therefore absolute values are compared
+            self.assertEqual(abs(actual_v), abs(expected_v))
+
+        shapes = (0, 3, 5)
+        batches = ((), (3, ), (2, 2))
+        for shape, batch in itertools.product(shapes, batches):
+            run_test(shape, batch)
+
+        # eigh requires a square matrix
+        t = torch.randn(2, 3, device=device, dtype=dtype)
+        with self.assertRaises(RuntimeError):
+            torch.linalg.eigh(t)
+
+    @skipCUDAIfNoMagma
+    @skipCPUIfNoLapack
+    @unittest.skipIf(not TEST_NUMPY, "NumPy not found")
+    @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128)
+    def test_eigvalsh(self, device, dtype):
+        from torch.testing._internal.common_utils import random_hermitian_matrix
+
+        def run_test(shape, batch):
+            matrix = random_hermitian_matrix(shape, *batch, dtype=dtype, device=device)
+            expected_w = np.linalg.eigvalsh(matrix.cpu().numpy())
+            actual_w = torch.linalg.eigvalsh(matrix)
+            self.assertEqual(actual_w, expected_w)
+
+        shapes = (0, 3, 5)
+        batches = ((), (3, ), (2, 2))
+        for shape, batch in itertools.product(shapes, batches):
+            run_test(shape, batch)
+
+        # eigvalsh requires a square matrix
+        t = torch.randn(2, 3, device=device, dtype=dtype)
+        with self.assertRaises(RuntimeError):
+            torch.linalg.eigvalsh(t)
+
     # This test confirms that torch.linalg.norm's dtype argument works
     # as expected, according to the function's documentation
     @skipCUDAIfNoMagma

torch/linalg/__init__.py

@@ -3,6 +3,8 @@
 import torch
 from torch._C import _add_docstr, _linalg  # type: ignore
 
+import functools
+
 Tensor = torch.Tensor
 
 # Note: This not only adds doc strings for functions in the linalg namespace, but
@@ -140,3 +142,93 @@
     >>> LA.norm(m[0, :, :]), LA.norm(m[1, :, :])
     (tensor(3.7417), tensor(11.2250))
 """)
+
+_add_docstr(_linalg.linalg_eigh, r"""
+linalg.eigh(input, UPLO='L') -> tuple(Tensor, Tensor)
+
+This function returns eigenvalues and eigenvectors
+of a complex Hermitian (conjugate symmetric) or real symmetric matrix :attr:`input`
+represented by a namedtuple (eigenvalues, eigenvectors).
+
+This function calculates all eigenvalues (and vectors) of :attr:`input`
+such that :math:`\text{input} = V \text{diag}(e) V^H`.
+
+Since the input matrix :attr:`input` is supposed to be Hermitian,
+only the lower triangular portion is used by default
+and the imaginary part of the diagonal will always be treated as zero.
+
+.. note:: The eigenvalues of real symmetric or complex Hermitian matrices are always real.
+
+.. note:: The eigenvalues/eigenvectors are computed using LAPACK routines ``_syevd``, ``_heevd``.
+
+Args:
+    input (Tensor): the input tensor of size :math:`(*, n, n)` where `*` is zero or more
+                    batch dimensions consisting of Hermitian matrices.
+    UPLO ('L', 'U', optional): controls whether to consider upper-triangular or lower-triangular part.
+        Default: ``'L'``
+
+Returns:
+    (Tensor, Tensor): A namedtuple (eigenvalues, eigenvectors) containing
+
+        - **eigenvalues** (*Tensor*): Shape :math:`(*, m)`.
+            The eigenvalues in ascending order, each repeated according to its multiplicity.
+        - **eigenvectors** (*Tensor*): Shape :math:`(*, m, m)`.
+            The orthonormal eigenvectors of the ``input``.
+
+Examples::
+
+    >>> import torch
+    >>> a = torch.randn(2, 2, dtype=torch.complex128)
+    >>> a = a + a.t().conj()  # To make a Hermitian
+    >>> a
+    tensor([[2.9228+0.0000j, 0.2029-0.0862j],
+            [0.2029+0.0862j, 0.3464+0.0000j]], dtype=torch.complex128)
+    >>> w, v = torch.linalg.eigh(a)
+    >>> w
+    tensor([0.3277, 2.9415], dtype=torch.float64)
+    >>> v
+    tensor([[-0.0846+-0.0000j, -0.9964+0.0000j],
+            [ 0.9170+0.3898j, -0.0779-0.0331j]], dtype=torch.complex128)
+    >>> torch.allclose(torch.matmul(v, torch.matmul(w.to(v.dtype).diag_embed(), v.transpose(-2, -1).conj())), a)
+    True
+""")
+
+_add_docstr(_linalg.linalg_eigvalsh, r"""
+linalg.eigvalsh(input, UPLO='L') -> Tensor
+
+This function returns eigenvalues of a complex Hermitian (conjugate symmetric)
+or real symmetric matrix :attr:`input`.
+
+.. note:: The eigenvalues of real symmetric or complex Hermitian matrices are always real.
+
+.. note:: The eigenvalues are computed using LAPACK routines ``_syevd``, ``_heevd``.
+
+Args:
+    input (Tensor): the input tensor of size :math:`(*, n, n)` where `*` is zero or more
+                    batch dimensions consisting of Hermitian matrices.
+    UPLO ('L', 'U', optional): controls whether to consider upper-triangular or lower-triangular part.
+        Default: ``'L'``
+
+Returns:
+    (Tensor): Shape :math:`(*, m)`. The eigenvalues in ascending order, each repeated according to its multiplicity.
+
+Examples::
+
+    >>> import torch
+    >>> a = torch.randn(2, 2, dtype=torch.complex128)
+    >>> a = a + a.t().conj()  # To make a Hermitian
+    >>> a
+    tensor([[2.9228+0.0000j, 0.2029-0.0862j],
+            [0.2029+0.0862j, 0.3464+0.0000j]], dtype=torch.complex128)
+    >>> w = torch.linalg.eigvalsh(a)
+    >>> w
+    tensor([0.3277, 2.9415], dtype=torch.float64)
+""")
+
+@functools.wraps(_linalg.linalg_eigh)
+def eigh(a, UPLO="L"):
+    return _linalg.linalg_eigh(a, UPLO)
+
+@functools.wraps(_linalg.linalg_eigh)
+def eigvalsh(a, UPLO="L"):
+    return _linalg.linalg_eigvalsh(a, UPLO)