Merge pull request #53 from rochisha0/add-sqrtm

Add dispatcher for sqrtm
qutip · Jun 12, 2024 · db0fb39 · db0fb39
2 parents fe51b69 + 6e20aeb
commit db0fb39
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/src/qutip_jax/unary.py b/src/qutip_jax/unary.py
@@ -18,6 +18,7 @@
     "conj_jaxdia",
     "inv_jaxarray",
     "expm_jaxarray",
+    "sqrtm_jaxarray",
     "project_jaxarray",
 ]
 
@@ -113,6 +114,12 @@ def inv_jaxarray(matrix):
     return JaxArray._fast_constructor(linalg.inv(matrix._jxa), matrix.shape)
 
 
+def sqrtm_jaxarray(matrix):
+    """Matrix square root `sqrt(A)` for a matrix `A`."""
+    _check_square_shape(matrix)
+    return JaxArray._fast_constructor(linalg.sqrtm(matrix._jxa), matrix.shape)
+
+
 @jit
 def _project_ket(array):
     return array @ array.T.conj()
@@ -185,6 +192,11 @@ def project_jaxarray(state):
 )
 
 
+qutip.data.sqrtm.add_specialisations(
+    [(JaxArray, JaxArray, sqrtm_jaxarray),]
+)
+
+
 qutip.data.project.add_specialisations(
     [
         (JaxArray, JaxArray, project_jaxarray),

diff --git a/tests/test_unary.py b/tests/test_unary.py
@@ -70,6 +70,12 @@ class TestInv(testing.TestInv):
     specialisations = [pytest.param(_inv_jax, JaxArray, JaxArray)]
 
 
+class TestSqrtm(testing.TestSqrtm):
+    specialisations = [
+        pytest.param(qutip_jax.sqrtm_jaxarray, JaxArray, JaxArray)
+    ]
+
+
 class TestProject(testing.TestProject):
     specialisations = [
         pytest.param(qutip_jax.project_jaxarray, JaxArray, JaxArray)