Implemented Cholesky decomposition to numpy and tensorflow backends

tensornetwork/backends/numpy/decompositions.py

@@ -122,3 +122,22 @@ def rq(
   r = np.reshape(r, list(left_dims) + [center_dim])
   q = np.reshape(q, [center_dim] + list(right_dims))
   return r, q
+
+
+def cholesky(
+    np: Any,
+    tensor: Tensor,
+    pivot_axis: int
+) -> Tuple[Tensor, Tensor]:
+  """
+  Computes the Cholesky decomposition of a tensor
+
+  See tensornetwork.backends.tensorflow.decompositions for details.
+  """
+  left_dims = np.shape(tensor)[:pivot_axis]
+  right_dims = np.shape(tensor)[pivot_axis:]
+  tensor = np.reshape(tensor,
+                      [np.reduce_prod(left_dims),
+                       np.reduce_prod(right_dims)])
+  L = np.linalg.cholesky(tensor)
+  return L

tensornetwork/backends/numpy/decompositions_test.py

@@ -16,7 +16,9 @@
 import numpy as np
 import tensorflow as tf
 from tensornetwork.backends.numpy import decompositions
+import pytest
 
+np_dtypes = [np.float64, np.complex128]
 
 class DecompositionsTest(tf.test.TestCase):
 
@@ -52,6 +54,13 @@ def test_qr(self):
       q, r = decompositions.qr(np, random_matrix, 1, non_negative_diagonal)
       self.assertAllClose(q.dot(r), random_matrix)
 
+  def test_cholesky(self):
+    #Assured positive-definite hermitian matrixs
+    random_matrix = np.random.rand(10, 10)
+    random_matrix = random_matrix @ random_matrix.T.conj()
+    L = decompositions.cholesky(tf, random_matrix, 1)
+    self.assertAllClose(np.linalg.cholesky(random_matrix), L)
+
   def test_max_singular_values(self):
     random_matrix = np.random.rand(10, 10)
     unitary1, _, unitary2 = np.linalg.svd(random_matrix)

tensornetwork/backends/pytorch/decompositions.py

@@ -121,6 +121,23 @@ def svd(
   return u, s, vh, s_rest
 
 
+def cholesky(
+    torch: Any,
+    tensor: Tensor,
+    pivot_axis: int
+) -> Tuple[Tensor, Tensor]:
+  """
+  Computes the Cholesky decomposition of a tensor
+
+  See tensornetwork.backends.tensorflow.decompositions for details.
+  """
+  left_dims = list(tensor.shape)[:pivot_axis]
+  right_dims = list(tensor.shape)[pivot_axis:]
+
+  tensor = torch.reshape(tensor, (np.prod(left_dims), np.prod(right_dims)))
+  L = np.linalg.cholesky(tensor)
+  return L
+
 def qr(
     torch: Any,
     tensor: Tensor,

tensornetwork/backends/pytorch/decompositions_test.py

@@ -16,7 +16,9 @@
 import numpy as np
 import torch
 from tensornetwork.backends.pytorch import decompositions
+import pytest
 
+np_dtypes = [np.float64, np.complex128]
 
 def test_expected_shapes():
   val = torch.zeros((2, 3, 4, 5))
@@ -42,6 +44,14 @@ def test_expected_shapes_rq():
     assert r.shape == (2, 3, 6)
     assert q.shape == (6, 4, 5)
 
+# @pytest.mark.parametrize("dtype", np_dtypes)
+def test_cholesky():
+  #Assured positive-definite hermitian matrix
+  random_matrix = np.random.rand(10, 10)
+  random_matrix = random_matrix @ random_matrix.T.conj()
+  random_matrix = torch.from_numpy(random_matrix)
+  L = decompositions.cholesky(torch, random_matrix, 1)
+  np.testing.assert_allclose(torch.cholesky(random_matrix), L)
 
 def test_rq():
   random_matrix = torch.rand([10, 10], dtype=torch.float64)

tensornetwork/backends/tensorflow/decompositions.py

@@ -226,3 +226,23 @@ def rq(
   r = tf.reshape(r, tf.concat([left_dims, [center_dim]], axis=-1))
   q = tf.reshape(q, tf.concat([[center_dim], right_dims], axis=-1))
   return r, q
+
+
+def cholesky(
+    tf: Any,
+    tensor: Tensor,
+    pivot_axis: int
+) -> Tuple[Tensor, Tensor]:
+  """ Computes de cholesky decomposition of a tensor.
+
+  Returns the Cholesky decomposition of a tensor which we treat as a 
+  square matrix
+  """
+  left_dims = tf.shape(tensor)[:pivot_axis]
+  right_dims = tf.shape(tensor)[pivot_axis:]
+  tensor = tf.reshape(tensor,
+                      [tf.reduce_prod(left_dims),
+                       tf.reduce_prod(right_dims)])
+  L = tf.linalg.cholesky(tensor)
+  return L
+

tensornetwork/backends/tensorflow/decompositions_test.py

@@ -16,7 +16,9 @@
 import numpy as np
 import tensorflow as tf
 from tensornetwork.backends.tensorflow import decompositions
+import pytest
 
+np_dtypes = [np.float64, np.complex128]
 
 class DecompositionsTest(tf.test.TestCase):
 
@@ -54,6 +56,14 @@ def test_qr(self):
       q, r = decompositions.qr(tf, random_matrix, 1, non_negative_diagonal)
       self.assertAllClose(tf.tensordot(q, r, ([1], [0])), random_matrix)
 
+  # @pytest.mark.parametrize("dtype", np_dtypes)
+  def test_cholesky(self):
+    #Assured positive-definite hermitian matrix
+    random_matrix = np.random.rand(10, 10)
+    random_matrix = random_matrix @ random_matrix.T.conj()
+    L = decompositions.cholesky(tf, random_matrix, 1)
+    self.assertAllClose(np.linalg.cholesky(random_matrix), L)
+
   def test_rq_defun(self):
     random_matrix = np.random.rand(10, 10)
     for non_negative_diagonal in [True, False]: