From a446a62808f88b004932e15a890d26afbedd6e4e Mon Sep 17 00:00:00 2001
From: abhinavd <abhinav.deshpande@ibm.com>
Date: Thu, 16 Jan 2025 14:55:22 -0800
Subject: [PATCH] Matrix square root for a not-perfectly PSD matrix

---
 tensorcircuit/backends/abstract_backend.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tensorcircuit/backends/abstract_backend.py b/tensorcircuit/backends/abstract_backend.py
index 4a7ccecd..1edf25ff 100644
--- a/tensorcircuit/backends/abstract_backend.py
+++ b/tensorcircuit/backends/abstract_backend.py
@@ -60,6 +60,21 @@ def sqrtmh(self: Any, a: Tensor) -> Tensor:
         e = self.sqrt(e)
         return v @ self.diagflat(e) @ self.adjoint(v)
 
+    def sqrtmhpos(self: Any, a: Tensor) -> Tensor:
+        """
+        Return the sqrtm of a known-to-be PSD Hermitian matrix ``a``.
+
+        :param a: tensor in matrix form
+        :type a: Tensor
+        :return: sqrtm of ``a`` after setting the negative eigenvalues (if they exist) to zero
+        :rtype: Tensor
+        """
+        # maybe friendly for AD and also cosidering that several backend has no support for native sqrtm
+        e, v = self.eigh(a)
+        e = self.relu(e)
+        e = self.sqrt(e)
+        return v @ self.diagflat(e) @ self.adjoint(v)
+
     def eigvalsh(self: Any, a: Tensor) -> Tensor:
         """
         Get the eigenvalues of matrix ``a``.