Merge pull request #110 from adtzlr/add-sqrtm

Add matrix square root `math.linalg.sqrtm()`

src/tensortrax/__about__.py

@@ -2,4 +2,4 @@
 tensorTRAX: Math on (Hyper-Dual) Tensors with Trailing Axes.
 """
 
-__version__ = "0.21.3"
+__version__ = "0.22.0"

src/tensortrax/math/linalg/__init__.py

@@ -5,6 +5,17 @@
 from ._linalg_array import det as _det
 from ._linalg_array import inv as _inv
 from ._linalg_array import pinv as _pinv
-from ._linalg_tensor import det, eigh, eigvalsh, expm, inv, pinv
+from ._linalg_tensor import det, eigh, eigvalsh, expm, inv, pinv, sqrtm
 
-__all__ = ["_det", "_inv", "_pinv", "det", "eigh", "eigvalsh", "expm", "inv", "pinv"]
+__all__ = [
+    "_det",
+    "_inv",
+    "_pinv",
+    "det",
+    "eigh",
+    "eigvalsh",
+    "expm",
+    "inv",
+    "pinv",
+    "sqrtm",
+]

src/tensortrax/math/linalg/_linalg_tensor.py

@@ -6,7 +6,7 @@
 import numpy as np
 
 from ..._tensor import Tensor, Δ, Δδ, einsum, f, matmul, δ
-from .._math_tensor import exp, sum, transpose
+from .._math_tensor import exp, sqrt, sum, transpose
 from ..special._special_tensor import ddot
 from . import _linalg_array as linalg
 
@@ -169,3 +169,9 @@ def expm(A):
     "Compute the matrix exponential of a symmetric array."
     λ, M = eigh(A)
     return einsum("a...,aij...->ij...", exp(λ), M)
+
+
+def sqrtm(A):
+    "Compute the matrix square root of a symmetric array."
+    λ, M = eigh(A)
+    return einsum("a...,aij...->ij...", sqrt(λ), M)

tests/test_math.py

@@ -89,6 +89,7 @@ def test_math():
     assert tm.linalg.eigvalsh(T).shape == (3,)
 
     assert tm.linalg.expm(T).shape == (3, 3)
+    assert tm.linalg.sqrtm(T).shape == (3, 3)
 
     assert tm.base.cross(F, F).shape == F.shape
     assert tm.base.eye(F).shape == F.shape