Implemented Cholesky decomposition to numpy and tensorflow backends

tensornetwork/backends/numpy/decompositions.py

@@ -122,3 +122,28 @@ 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(
+    tf: Any,
+    tensor: Tensor,
+    pivot_axis: int,
+    non_negative_diagonal: bool
+) -> Tuple[Tensor, Tensor]:
+  """
+  Computes the Cholesky decomposition of a tensor
+
+  See tensornetwork.backends.tensorflow.decompositions for details.
+  """
+  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)
+  if non_negative_diagonal:
+    phases = tf.math.sign(tf.linalg.diag_part(L))
+    L = phases[:, None] * L
+  center_dim = tf.shape(L)[1]
+  L = tf.reshape(L, tf.concat([left_dims, [center_dim]], axis=-1))
+  return L

tensornetwork/backends/numpy/decompositions_test.py

@@ -52,6 +52,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 = [[1.0, 0], [0, 1.0]]
+    for non_negative_diagonal in [True, False]:
+      L = decompositions.cholesky(tf, random_matrix, 1, non_negative_diagonal)
+      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,29 @@ def svd(
   return u, s, vh, s_rest
 
 
+def cholesky(
+    torch: Any,
+    tensor: Tensor,
+    pivot_axis: int,
+    non_negative_diagonal: bool = False
+) -> 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 = torch.reshape(tensor, (np.prod(left_dims), np.prod(right_dims)))
+  L = torch.cholesky(tensor)
+  if non_negative_diagonal:
+    phases = torch.sign(torch.diagonal(L))
+    L = phases[:, None] * L
+  center_dim = L.shape[1]
+  L = torch.reshape(L, list(left_dims) + [center_dim])
+  return L
+
 def qr(
     torch: Any,
     tensor: Tensor,

tensornetwork/backends/pytorch/decompositions_test.py

@@ -42,6 +42,13 @@ def test_expected_shapes_rq():
     assert r.shape == (2, 3, 6)
     assert q.shape == (6, 4, 5)
 
+def test_cholesky():
+  #Assured positive-definite hermitian matrix
+  random_matrix = np.array([[1.0, 0], [0, 1.0]])
+  random_matrix = torch.from_numpy(random_matrix)
+  for non_negative_diagonal in [True, False]:
+    L = decompositions.cholesky(torch, random_matrix, 1, non_negative_diagonal)
+    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,28 @@ 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,
+    non_negative_diagonal: bool
+) -> 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)
+  if non_negative_diagonal:
+    phases = tf.math.sign(tf.linalg.diag_part(L))
+    L = phases[:, None] * L
+  center_dim = tf.shape(L)[1]
+  L = tf.reshape(L, tf.concat([left_dims, [center_dim]], axis=-1))
+  return L

tensornetwork/backends/tensorflow/decompositions_test.py

@@ -54,6 +54,13 @@ 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)
 
+  def test_cholesky(self):
+    #Assured positive-definite hermitian matrix
+    random_matrix = [[1.0, 0], [0, 1.0]]
+    for non_negative_diagonal in [True, False]:
+      L = decompositions.cholesky(tf, random_matrix, 1, non_negative_diagonal)
+      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]: