From 016d654f2dc8b61b15457f62eff1719c4881e6d2 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Fri, 4 Dec 2020 18:14:57 -0700
Subject: [PATCH 01/21] Added power to Jax backend and test

---
 tensornetwork/backends/jax/jax_backend.py      |  3 +++
 tensornetwork/backends/jax/jax_backend_test.py | 15 +++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index d7417e8f3..038960ce3 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -874,3 +874,6 @@ def sign(self, tensor: Tensor) -> Tensor:
 
   def item(self, tensor):
     return tensor.item()
+
+  def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
+    return jnp.power(a,b)
\ No newline at end of file
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
index 13c74247c..cc3246b57 100644
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -974,6 +974,18 @@ def matvec_jax(vector, matrix):
       num_krylov_vecs=100,
       tol=0.0001)
 
+def test_power(dtype):
+  np.random.seed(10)
+  backend = jax_backend.JaxBackend()
+  tensor = np.random.rand(2, 3, 4)
+  a = backend.convert_to_tensor(tensor)
+  actual = backend.power(a, axis=(1, 2))
+  expected = np.power(a, axis=(1, 2))
+  np.testing.assert_allclose(expected, actual)
+
+  actual = backend.power(a, axis=(1, 2), keepdims=True)
+  expected = np.power(a, axis=(1, 2), keepdims=True)
+  np.testing.assert_allclose(expected, actual)
 
 def test_sum():
   np.random.seed(10)
@@ -1240,3 +1252,6 @@ def test_item(dtype):
   backend = jax_backend.JaxBackend()
   tensor = backend.randn((1,), dtype=dtype, seed=10)
   assert backend.item(tensor) == tensor.item()
+
+
+

From 62a06c0688a04a45fa80945c08f311eace398bda Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Fri, 4 Dec 2020 19:18:27 -0700
Subject: [PATCH 02/21] Added power function to the Jax backend and test.

---
 tensornetwork/backends/jax/jax_backend.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index 038960ce3..aa488bee8 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -876,4 +876,15 @@ def item(self, tensor):
     return tensor.item()
 
   def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
-    return jnp.power(a,b)
\ No newline at end of file
+    """
+    Returns power of tensor a to the value of b.
+    In the case b is a tensor, then the power is by element
+      with a as the base and b as the exponent.
+    In the case b is a scalar, then the power of each value in a
+      is raised to the exponent of b.
+
+    Args:
+      a: The tensor that contains the base.
+      b: The tensor that contains the exponent or a single scalar.
+    """
+    return jnp.power(a,b)

From 43952dde29765a0dfcd73c97b28914a55d9c7ad3 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Fri, 4 Dec 2020 20:13:21 -0700
Subject: [PATCH 03/21] Fixed white space error.

---
 tensornetwork/backends/jax/jax_backend.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index aa488bee8..68c4d9540 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -887,4 +887,4 @@ def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
       a: The tensor that contains the base.
       b: The tensor that contains the exponent or a single scalar.
     """
-    return jnp.power(a,b)
+    return jnp.power(a, b)

From b0f579e4634e3636793dd8b3f933ad14eb19b123 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sat, 5 Dec 2020 00:44:28 -0700
Subject: [PATCH 04/21] Re-made the power test function.

---
 .../backends/jax/jax_backend_test.py          | 29 +++++++++----------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
index cc3246b57..3c8926f64 100644
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -974,19 +974,6 @@ def matvec_jax(vector, matrix):
       num_krylov_vecs=100,
       tol=0.0001)
 
-def test_power(dtype):
-  np.random.seed(10)
-  backend = jax_backend.JaxBackend()
-  tensor = np.random.rand(2, 3, 4)
-  a = backend.convert_to_tensor(tensor)
-  actual = backend.power(a, axis=(1, 2))
-  expected = np.power(a, axis=(1, 2))
-  np.testing.assert_allclose(expected, actual)
-
-  actual = backend.power(a, axis=(1, 2), keepdims=True)
-  expected = np.power(a, axis=(1, 2), keepdims=True)
-  np.testing.assert_allclose(expected, actual)
-
 def test_sum():
   np.random.seed(10)
   backend = jax_backend.JaxBackend()
@@ -1253,5 +1240,17 @@ def test_item(dtype):
   tensor = backend.randn((1,), dtype=dtype, seed=10)
   assert backend.item(tensor) == tensor.item()
 
-
-
+@pytest.mark.parametrize("dtype", tf_dtypes)
+def test_power(dtype):
+  shape = (4, 3, 2)
+  backend = jax_backend.JaxBackend()
+  base_tensor = backend.randn(shape, dtype=dtype, seed=10)
+  power_tensor = backend.randn(shape, dtype=dtype, seed=10)
+  actual = backend.power(base_tensor, power_tensor)
+  expected = tf.math.pow(base_tensor, power_tensor)
+  np.testing.assert_allclose(expected, actual)
+  
+  power = np.random.rand(1)[0]
+  actual = backend.power(base_tensor, power)
+  expected = tf.math.pow(base_tensor, power)
+  np.testing.assert_allclose(expected, actual)
\ No newline at end of file

From e9f9869942b9ceb58caf5552e6e7290be0a958a7 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sat, 5 Dec 2020 00:49:59 -0700
Subject: [PATCH 05/21] Fixed typo

---
 tensornetwork/backends/jax/jax_backend_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
index 3c8926f64..02b41c05e 100644
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -1240,7 +1240,7 @@ def test_item(dtype):
   tensor = backend.randn((1,), dtype=dtype, seed=10)
   assert backend.item(tensor) == tensor.item()
 
-@pytest.mark.parametrize("dtype", tf_dtypes)
+@pytest.mark.parametrize("dtype", np_dtypes)
 def test_power(dtype):
   shape = (4, 3, 2)
   backend = jax_backend.JaxBackend()
@@ -1249,7 +1249,7 @@ def test_power(dtype):
   actual = backend.power(base_tensor, power_tensor)
   expected = tf.math.pow(base_tensor, power_tensor)
   np.testing.assert_allclose(expected, actual)
-  
+
   power = np.random.rand(1)[0]
   actual = backend.power(base_tensor, power)
   expected = tf.math.pow(base_tensor, power)

From 9e16a9faacbfc318a497f8dadf130b4ef0ca32c1 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sat, 5 Dec 2020 01:30:23 -0700
Subject: [PATCH 06/21] Testing out numpy square.

---
 tensornetwork/backends/jax/jax_backend.py      | 2 +-
 tensornetwork/backends/jax/jax_backend_test.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index 68c4d9540..f3d7c04a2 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -887,4 +887,4 @@ def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
       a: The tensor that contains the base.
       b: The tensor that contains the exponent or a single scalar.
     """
-    return jnp.power(a, b)
+    return jnp.square(a, b)
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
index 02b41c05e..6050c6271 100644
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -1253,4 +1253,5 @@ def test_power(dtype):
   power = np.random.rand(1)[0]
   actual = backend.power(base_tensor, power)
   expected = tf.math.pow(base_tensor, power)
-  np.testing.assert_allclose(expected, actual)
\ No newline at end of file
+  np.testing.assert_allclose(expected, actual)
+  
\ No newline at end of file

From 81b4d8b6faeeb251e3ec8ef040c2ba174d074d4d Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sat, 5 Dec 2020 02:09:54 -0700
Subject: [PATCH 07/21] Fixed issues with the assertion in test, should work
 now.

---
 tensornetwork/backends/jax/jax_backend.py      | 4 ++--
 tensornetwork/backends/jax/jax_backend_test.py | 5 ++---
 2 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index f3d7c04a2..2ef092eda 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -875,7 +875,7 @@ def sign(self, tensor: Tensor) -> Tensor:
   def item(self, tensor):
     return tensor.item()
 
-  def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
+  def power(self, a: Tensor, b: Union[Tensor, int]) -> Tensor:
     """
     Returns power of tensor a to the value of b.
     In the case b is a tensor, then the power is by element
@@ -887,4 +887,4 @@ def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
       a: The tensor that contains the base.
       b: The tensor that contains the exponent or a single scalar.
     """
-    return jnp.square(a, b)
+    return jnp.power(a, b)
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
index 6050c6271..87b270e02 100644
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -1247,11 +1247,10 @@ def test_power(dtype):
   base_tensor = backend.randn(shape, dtype=dtype, seed=10)
   power_tensor = backend.randn(shape, dtype=dtype, seed=10)
   actual = backend.power(base_tensor, power_tensor)
-  expected = tf.math.pow(base_tensor, power_tensor)
+  expected = jax.numpy.power(base_tensor, power_tensor)
   np.testing.assert_allclose(expected, actual)
 
   power = np.random.rand(1)[0]
   actual = backend.power(base_tensor, power)
-  expected = tf.math.pow(base_tensor, power)
+  expected = jax.numpy.power(base_tensor, power)
   np.testing.assert_allclose(expected, actual)
-  
\ No newline at end of file

From 195f69198d4cc1a626ca43bba07bab05ece994fd Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sat, 5 Dec 2020 23:57:51 -0700
Subject: [PATCH 08/21] Added NotImplementedError function and it's respective
 test for Cholskey decomposition.

---
 tensornetwork/backends/abstract_backend.py | 4 ++++
 tensornetwork/backends/backend_test.py     | 6 ++++++
 2 files changed, 10 insertions(+)

diff --git a/tensornetwork/backends/abstract_backend.py b/tensornetwork/backends/abstract_backend.py
index 4e57d366e..e72d5a14f 100644
--- a/tensornetwork/backends/abstract_backend.py
+++ b/tensornetwork/backends/abstract_backend.py
@@ -1022,3 +1022,7 @@ def item(self, tensor) -> Union[float, int, complex]:
       The value in tensor.
     """
     raise NotImplementedError("Backend {self.name} has not implemented item")
+
+  def chsky(self, tensor: Tensor, pivot_axis: int = -1, non_negative_diagonal: bool = False) -> Tuple[Tensor, Tensor]:
+    """Computes the Cholskey decomposition of a tensor."""
+    raise NotImplementedError("Backend '{}' has not implemented chsky.".format(self.name))
diff --git a/tensornetwork/backends/backend_test.py b/tensornetwork/backends/backend_test.py
index 6230361a1..2aef64732 100644
--- a/tensornetwork/backends/backend_test.py
+++ b/tensornetwork/backends/backend_test.py
@@ -193,6 +193,12 @@ def test_abstract_backend_qr_decompositon_not_implemented():
     backend.qr(np.ones((2, 2)), 0)
 
 
+def test_abstract_backend_chsky_decompositon_not_implemented():
+  backend = AbstractBackend()
+  with pytest.raises(NotImplementedError):
+    backend.chsky(np.ones((2, 2)), 0)
+
+
 def test_abstract_backend_rq_decompositon_not_implemented():
   backend = AbstractBackend()
   with pytest.raises(NotImplementedError):

From 737f07149cbb3cc09399d30b3035f02171954a56 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sun, 6 Dec 2020 00:23:20 -0700
Subject: [PATCH 09/21] Fixed line too long error

---
 tensornetwork/backends/abstract_backend.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensornetwork/backends/abstract_backend.py b/tensornetwork/backends/abstract_backend.py
index e72d5a14f..31ff08921 100644
--- a/tensornetwork/backends/abstract_backend.py
+++ b/tensornetwork/backends/abstract_backend.py
@@ -1023,6 +1023,9 @@ def item(self, tensor) -> Union[float, int, complex]:
     """
     raise NotImplementedError("Backend {self.name} has not implemented item")
 
-  def chsky(self, tensor: Tensor, pivot_axis: int = -1, non_negative_diagonal: bool = False) -> Tuple[Tensor, Tensor]:
+  def chsky(self, tensor: Tensor, 
+            pivot_axis: int = -1, 
+            non_negative_diagonal: bool = False) -> 
+            Tuple[Tensor, Tensor]:
     """Computes the Cholskey decomposition of a tensor."""
     raise NotImplementedError("Backend '{}' has not implemented chsky.".format(self.name))

From 0fbe8fb3cb5a94b219280f17e65abbfebccdb83f Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sun, 6 Dec 2020 00:59:34 -0700
Subject: [PATCH 10/21] Removing changes for different branches.

---
 tensornetwork/backends/jax/jax_backend.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index 2ef092eda..d444653b1 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -877,7 +877,7 @@ def item(self, tensor):
 
   def power(self, a: Tensor, b: Union[Tensor, int]) -> Tensor:
     """
-    Returns power of tensor a to the value of b.
+    Returns the power of tensor a to the value of b.
     In the case b is a tensor, then the power is by element
       with a as the base and b as the exponent.
     In the case b is a scalar, then the power of each value in a

From 8c376566cc1ab9e8151bed64edbd2fc538a802b2 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sun, 6 Dec 2020 01:38:44 -0700
Subject: [PATCH 11/21] Revert "Fixed issues with the assertion in test, should
 work now."

This reverts commit 81b4d8b6faeeb251e3ec8ef040c2ba174d074d4d.
---
 tensornetwork/backends/jax/jax_backend.py      | 4 ++--
 tensornetwork/backends/jax/jax_backend_test.py | 5 +++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index d444653b1..3b6ed659a 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -875,7 +875,7 @@ def sign(self, tensor: Tensor) -> Tensor:
   def item(self, tensor):
     return tensor.item()
 
-  def power(self, a: Tensor, b: Union[Tensor, int]) -> Tensor:
+  def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
     """
     Returns the power of tensor a to the value of b.
     In the case b is a tensor, then the power is by element
@@ -887,4 +887,4 @@ def power(self, a: Tensor, b: Union[Tensor, int]) -> Tensor:
       a: The tensor that contains the base.
       b: The tensor that contains the exponent or a single scalar.
     """
-    return jnp.power(a, b)
+    return jnp.square(a, b)
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
index 87b270e02..6050c6271 100644
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -1247,10 +1247,11 @@ def test_power(dtype):
   base_tensor = backend.randn(shape, dtype=dtype, seed=10)
   power_tensor = backend.randn(shape, dtype=dtype, seed=10)
   actual = backend.power(base_tensor, power_tensor)
-  expected = jax.numpy.power(base_tensor, power_tensor)
+  expected = tf.math.pow(base_tensor, power_tensor)
   np.testing.assert_allclose(expected, actual)
 
   power = np.random.rand(1)[0]
   actual = backend.power(base_tensor, power)
-  expected = jax.numpy.power(base_tensor, power)
+  expected = tf.math.pow(base_tensor, power)
   np.testing.assert_allclose(expected, actual)
+  
\ No newline at end of file

From a89496018d68557c806ffc0f898bb964a0c3537f Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sun, 6 Dec 2020 11:34:26 -0700
Subject: [PATCH 12/21] Fixed a few pylint issues, should work now

---
 tensornetwork/backends/abstract_backend.py | 13 +++++++------
 tensornetwork/backends/backend_test.py     |  4 ++--
 tensornetwork/backends/jax/jax_backend.py  |  2 +-
 3 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/tensornetwork/backends/abstract_backend.py b/tensornetwork/backends/abstract_backend.py
index 31ff08921..f439c3113 100644
--- a/tensornetwork/backends/abstract_backend.py
+++ b/tensornetwork/backends/abstract_backend.py
@@ -1023,9 +1023,10 @@ def item(self, tensor) -> Union[float, int, complex]:
     """
     raise NotImplementedError("Backend {self.name} has not implemented item")
 
-  def chsky(self, tensor: Tensor, 
-            pivot_axis: int = -1, 
-            non_negative_diagonal: bool = False) -> 
-            Tuple[Tensor, Tensor]:
-    """Computes the Cholskey decomposition of a tensor."""
-    raise NotImplementedError("Backend '{}' has not implemented chsky.".format(self.name))
+  def cholesky_decomposition(self, 
+                             tensor: Tensor,
+                             pivot_axis: int = -1,
+                             non_negative_diagonal: bool = False) -> \
+                             Tuple[Tensor, Tensor]:
+    raise NotImplementedError(
+        f"Backend {self.name} han not implemented cholesky_decomposition.")
diff --git a/tensornetwork/backends/backend_test.py b/tensornetwork/backends/backend_test.py
index 2aef64732..d368173f3 100644
--- a/tensornetwork/backends/backend_test.py
+++ b/tensornetwork/backends/backend_test.py
@@ -193,10 +193,10 @@ def test_abstract_backend_qr_decompositon_not_implemented():
     backend.qr(np.ones((2, 2)), 0)
 
 
-def test_abstract_backend_chsky_decompositon_not_implemented():
+def test_abstract_backend_cholesky_decompositon_not_implemented():
   backend = AbstractBackend()
   with pytest.raises(NotImplementedError):
-    backend.chsky(np.ones((2, 2)), 0)
+    backend.cholesky_decomposition(np.ones((2, 2)), 0)
 
 
 def test_abstract_backend_rq_decompositon_not_implemented():
diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index 3b6ed659a..9f5eb9843 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -887,4 +887,4 @@ def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
       a: The tensor that contains the base.
       b: The tensor that contains the exponent or a single scalar.
     """
-    return jnp.square(a, b)
+    return jnp.power(a, b)

From 3ba747f0264d0eba96cee3c5864032982787aca7 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Sun, 6 Dec 2020 12:41:10 -0700
Subject: [PATCH 13/21] Fixed typo in the error message

---
 tensornetwork/backends/abstract_backend.py     |  2 +-
 tensornetwork/backends/jax/jax_backend.py      | 14 --------------
 tensornetwork/backends/jax/jax_backend_test.py | 15 ---------------
 3 files changed, 1 insertion(+), 30 deletions(-)

diff --git a/tensornetwork/backends/abstract_backend.py b/tensornetwork/backends/abstract_backend.py
index f439c3113..632c05613 100644
--- a/tensornetwork/backends/abstract_backend.py
+++ b/tensornetwork/backends/abstract_backend.py
@@ -1029,4 +1029,4 @@ def cholesky_decomposition(self,
                              non_negative_diagonal: bool = False) -> \
                              Tuple[Tensor, Tensor]:
     raise NotImplementedError(
-        f"Backend {self.name} han not implemented cholesky_decomposition.")
+        f"Backend {self.name} has not implemented cholesky_decomposition.")
diff --git a/tensornetwork/backends/jax/jax_backend.py b/tensornetwork/backends/jax/jax_backend.py
index 9f5eb9843..d7417e8f3 100644
--- a/tensornetwork/backends/jax/jax_backend.py
+++ b/tensornetwork/backends/jax/jax_backend.py
@@ -874,17 +874,3 @@ def sign(self, tensor: Tensor) -> Tensor:
 
   def item(self, tensor):
     return tensor.item()
-
-  def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
-    """
-    Returns the power of tensor a to the value of b.
-    In the case b is a tensor, then the power is by element
-      with a as the base and b as the exponent.
-    In the case b is a scalar, then the power of each value in a
-      is raised to the exponent of b.
-
-    Args:
-      a: The tensor that contains the base.
-      b: The tensor that contains the exponent or a single scalar.
-    """
-    return jnp.power(a, b)
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
index 6050c6271..63e761820 100644
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -1239,19 +1239,4 @@ def test_item(dtype):
   backend = jax_backend.JaxBackend()
   tensor = backend.randn((1,), dtype=dtype, seed=10)
   assert backend.item(tensor) == tensor.item()
-
-@pytest.mark.parametrize("dtype", np_dtypes)
-def test_power(dtype):
-  shape = (4, 3, 2)
-  backend = jax_backend.JaxBackend()
-  base_tensor = backend.randn(shape, dtype=dtype, seed=10)
-  power_tensor = backend.randn(shape, dtype=dtype, seed=10)
-  actual = backend.power(base_tensor, power_tensor)
-  expected = tf.math.pow(base_tensor, power_tensor)
-  np.testing.assert_allclose(expected, actual)
-
-  power = np.random.rand(1)[0]
-  actual = backend.power(base_tensor, power)
-  expected = tf.math.pow(base_tensor, power)
-  np.testing.assert_allclose(expected, actual)
   
\ No newline at end of file

From f9fb40db4c9b5766cdb3ed180489986d27acce2e Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Mon, 7 Dec 2020 14:37:07 -0700
Subject: [PATCH 14/21] Renamed cholesky function, to a shorter name

---
 tensornetwork/backends/abstract_backend.py | 10 +++++-----
 tensornetwork/backends/backend_test.py     |  2 +-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensornetwork/backends/abstract_backend.py b/tensornetwork/backends/abstract_backend.py
index 632c05613..042b52995 100644
--- a/tensornetwork/backends/abstract_backend.py
+++ b/tensornetwork/backends/abstract_backend.py
@@ -1023,10 +1023,10 @@ def item(self, tensor) -> Union[float, int, complex]:
     """
     raise NotImplementedError("Backend {self.name} has not implemented item")
 
-  def cholesky_decomposition(self, 
-                             tensor: Tensor,
-                             pivot_axis: int = -1,
-                             non_negative_diagonal: bool = False) -> \
-                             Tuple[Tensor, Tensor]:
+  def cholesky(self, 
+               tensor: Tensor,
+               pivot_axis: int = -1,
+               non_negative_diagonal: bool = False) -> \
+               Tuple[Tensor, Tensor]:
     raise NotImplementedError(
         f"Backend {self.name} has not implemented cholesky_decomposition.")
diff --git a/tensornetwork/backends/backend_test.py b/tensornetwork/backends/backend_test.py
index d368173f3..41cf7d999 100644
--- a/tensornetwork/backends/backend_test.py
+++ b/tensornetwork/backends/backend_test.py
@@ -196,7 +196,7 @@ def test_abstract_backend_qr_decompositon_not_implemented():
 def test_abstract_backend_cholesky_decompositon_not_implemented():
   backend = AbstractBackend()
   with pytest.raises(NotImplementedError):
-    backend.cholesky_decomposition(np.ones((2, 2)), 0)
+    backend.cholesky(np.ones((2, 2)), 0)
 
 
 def test_abstract_backend_rq_decompositon_not_implemented():

From 0c3817dc455f99d158b5f13b226d295997da4137 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Mon, 7 Dec 2020 14:47:33 -0700
Subject: [PATCH 15/21] Renamed cholesky function from cholesky_decomposition
 to cholesky

---
 tensornetwork/backends/abstract_backend.py     | 2 +-
 tensornetwork/backends/jax/jax_backend_test.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensornetwork/backends/abstract_backend.py b/tensornetwork/backends/abstract_backend.py
index 042b52995..bcde899b1 100644
--- a/tensornetwork/backends/abstract_backend.py
+++ b/tensornetwork/backends/abstract_backend.py
@@ -1029,4 +1029,4 @@ def cholesky(self,
                non_negative_diagonal: bool = False) -> \
                Tuple[Tensor, Tensor]:
     raise NotImplementedError(
-        f"Backend {self.name} has not implemented cholesky_decomposition.")
+        f"Backend {self.name} has not implemented cholesky.")
diff --git a/tensornetwork/backends/jax/jax_backend_test.py b/tensornetwork/backends/jax/jax_backend_test.py
index 9911c3ab5..615bba2e6 100644
--- a/tensornetwork/backends/jax/jax_backend_test.py
+++ b/tensornetwork/backends/jax/jax_backend_test.py
@@ -1253,4 +1253,5 @@ def test_power(dtype):
   power = np.random.rand(1)[0]
   actual = backend.power(base_tensor, power)
   expected = jax.numpy.power(base_tensor, power)
-  np.testing.assert_allclose(expected, actual)
\ No newline at end of file
+  np.testing.assert_allclose(expected, actual)
+  
\ No newline at end of file

From aedfe2ace8792eb54ac849e185f2a9e1257931fb Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Wed, 23 Dec 2020 14:30:56 -0700
Subject: [PATCH 16/21] Implemented Cholesky to numpy and tensorflow with their
 respective tests

---
 .../backends/numpy/decompositions.py          | 19 +++++++++++++++++++
 .../backends/numpy/decompositions_test.py     |  6 ++++++
 2 files changed, 25 insertions(+)

diff --git a/tensornetwork/backends/numpy/decompositions.py b/tensornetwork/backends/numpy/decompositions.py
index 29feb63d0..129bc6e92 100644
--- a/tensornetwork/backends/numpy/decompositions.py
+++ b/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, # TODO: Typing
+    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,
+                      [numpy.prod(left_dims),
+                       numpy.prod(right_dims)])
+  L = np.linalg.cholesky(tensor)
+  return L
\ No newline at end of file
diff --git a/tensornetwork/backends/numpy/decompositions_test.py b/tensornetwork/backends/numpy/decompositions_test.py
index c03801975..813d5536d 100644
--- a/tensornetwork/backends/numpy/decompositions_test.py
+++ b/tensornetwork/backends/numpy/decompositions_test.py
@@ -52,6 +52,12 @@ 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):
+    random_matrix = np.random.rand(10, 10)
+    for non_negative_diagonal in [True, False]:
+      L = decompositions.cholesky(np, random_matrix, 1)
+      self.assertAllClose(np.cholesky(random_matrix), L)
+
   def test_max_singular_values(self):
     random_matrix = np.random.rand(10, 10)
     unitary1, _, unitary2 = np.linalg.svd(random_matrix)

From 13026b5c12d77deffc72175d2c6987141dca472e Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Wed, 23 Dec 2020 15:17:57 -0700
Subject: [PATCH 17/21] Fixed pylint issues

---
 tensornetwork/backends/numpy/decompositions_test.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensornetwork/backends/numpy/decompositions_test.py b/tensornetwork/backends/numpy/decompositions_test.py
index 813d5536d..51964e321 100644
--- a/tensornetwork/backends/numpy/decompositions_test.py
+++ b/tensornetwork/backends/numpy/decompositions_test.py
@@ -54,9 +54,8 @@ def test_qr(self):
 
   def test_cholesky(self):
     random_matrix = np.random.rand(10, 10)
-    for non_negative_diagonal in [True, False]:
-      L = decompositions.cholesky(np, random_matrix, 1)
-      self.assertAllClose(np.cholesky(random_matrix), L)
+    L = decompositions.cholesky(np, random_matrix, 1)
+    self.assertAllClose(np.cholesky(random_matrix), L)
 
   def test_max_singular_values(self):
     random_matrix = np.random.rand(10, 10)

From 11eba6af6e98119c75e295ee2717bd4aba5d8b76 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Fri, 25 Dec 2020 12:50:46 -0700
Subject: [PATCH 18/21] Implemented Cholesky decomposition to tensorflow, numpy
 & pytorch backends, fixing various pylint errors and adding test to each
 respective backend function

---
 .../backends/numpy/decompositions.py          | 24 +++++++++++-------
 .../backends/numpy/decompositions_test.py     |  7 +++---
 .../backends/pytorch/decompositions.py        | 23 +++++++++++++++++
 .../backends/pytorch/decompositions_test.py   |  7 ++++++
 .../backends/tensorflow/decompositions.py     | 25 +++++++++++++++++++
 .../tensorflow/decompositions_test.py         |  6 +++++
 6 files changed, 80 insertions(+), 12 deletions(-)

diff --git a/tensornetwork/backends/numpy/decompositions.py b/tensornetwork/backends/numpy/decompositions.py
index 129bc6e92..4f47526d0 100644
--- a/tensornetwork/backends/numpy/decompositions.py
+++ b/tensornetwork/backends/numpy/decompositions.py
@@ -125,19 +125,25 @@ def rq(
 
 
 def cholesky(
-    np, # TODO: Typing
+    tf: Any,
     tensor: Tensor,
-    pivot_axis: int
+    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 = np.shape(tensor)[:pivot_axis]
-  right_dims = np.shape(tensor)[pivot_axis:]
-  tensor = np.reshape(tensor,
-                      [numpy.prod(left_dims),
-                       numpy.prod(right_dims)])
-  L = np.linalg.cholesky(tensor)
-  return L
\ No newline at end of file
+  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
diff --git a/tensornetwork/backends/numpy/decompositions_test.py b/tensornetwork/backends/numpy/decompositions_test.py
index 51964e321..04c5ef210 100644
--- a/tensornetwork/backends/numpy/decompositions_test.py
+++ b/tensornetwork/backends/numpy/decompositions_test.py
@@ -53,9 +53,10 @@ def test_qr(self):
       self.assertAllClose(q.dot(r), random_matrix)
 
   def test_cholesky(self):
-    random_matrix = np.random.rand(10, 10)
-    L = decompositions.cholesky(np, random_matrix, 1)
-    self.assertAllClose(np.cholesky(random_matrix), L)
+    random_matrix = [[1.0, 0], [0, 1.0]] #Assured positive-definite hermitian matrix
+    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)
diff --git a/tensornetwork/backends/pytorch/decompositions.py b/tensornetwork/backends/pytorch/decompositions.py
index 773902ea0..535807e2d 100644
--- a/tensornetwork/backends/pytorch/decompositions.py
+++ b/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,
diff --git a/tensornetwork/backends/pytorch/decompositions_test.py b/tensornetwork/backends/pytorch/decompositions_test.py
index 908fdf0cc..bafcd2447 100644
--- a/tensornetwork/backends/pytorch/decompositions_test.py
+++ b/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)
diff --git a/tensornetwork/backends/tensorflow/decompositions.py b/tensornetwork/backends/tensorflow/decompositions.py
index dcc5bba40..2ce0aae93 100644
--- a/tensornetwork/backends/tensorflow/decompositions.py
+++ b/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
\ No newline at end of file
diff --git a/tensornetwork/backends/tensorflow/decompositions_test.py b/tensornetwork/backends/tensorflow/decompositions_test.py
index 05ebb3c17..b42c03c5f 100644
--- a/tensornetwork/backends/tensorflow/decompositions_test.py
+++ b/tensornetwork/backends/tensorflow/decompositions_test.py
@@ -54,6 +54,12 @@ 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):
+    random_matrix = [[1.0, 0], [0, 1.0]] #Assured positive-definite hermitian matrix
+    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]:

From 4da6eacc94940335949941f03803264d837f8aa1 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Fri, 25 Dec 2020 12:56:50 -0700
Subject: [PATCH 19/21] Fixed minor pylint issues

---
 tensornetwork/backends/numpy/decompositions_test.py      | 3 ++-
 tensornetwork/backends/tensorflow/decompositions.py      | 2 +-
 tensornetwork/backends/tensorflow/decompositions_test.py | 3 ++-
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensornetwork/backends/numpy/decompositions_test.py b/tensornetwork/backends/numpy/decompositions_test.py
index 04c5ef210..ac04afdcc 100644
--- a/tensornetwork/backends/numpy/decompositions_test.py
+++ b/tensornetwork/backends/numpy/decompositions_test.py
@@ -53,7 +53,8 @@ def test_qr(self):
       self.assertAllClose(q.dot(r), random_matrix)
 
   def test_cholesky(self):
-    random_matrix = [[1.0, 0], [0, 1.0]] #Assured positive-definite hermitian matrix
+    #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)
diff --git a/tensornetwork/backends/tensorflow/decompositions.py b/tensornetwork/backends/tensorflow/decompositions.py
index 2ce0aae93..4f92a2bbe 100644
--- a/tensornetwork/backends/tensorflow/decompositions.py
+++ b/tensornetwork/backends/tensorflow/decompositions.py
@@ -249,5 +249,5 @@ def cholesky(
     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))
+  L = tf.reshape(L, tf.concat([left_dims, [center_dim]], axis=-1))
   return L
\ No newline at end of file
diff --git a/tensornetwork/backends/tensorflow/decompositions_test.py b/tensornetwork/backends/tensorflow/decompositions_test.py
index b42c03c5f..87042264d 100644
--- a/tensornetwork/backends/tensorflow/decompositions_test.py
+++ b/tensornetwork/backends/tensorflow/decompositions_test.py
@@ -55,7 +55,8 @@ def test_qr(self):
       self.assertAllClose(tf.tensordot(q, r, ([1], [0])), random_matrix)
 
   def test_cholesky(self):
-    random_matrix = [[1.0, 0], [0, 1.0]] #Assured positive-definite hermitian matrix
+    #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)

From dab2060242743521b6b3fccb8f40704cdfc4e191 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Mon, 28 Dec 2020 11:23:36 -0700
Subject: [PATCH 20/21] requested changes made removed unnecessary arguments
 and better matrix

---
 .../backends/numpy/decompositions.py          |  22 +--
 .../backends/numpy/decompositions_test.py     |  10 +-
 .../backends/pytorch/decompositions.py        |  16 +-
 .../backends/pytorch/decompositions_test.py   | 175 +++++++++---------
 .../backends/tensorflow/decompositions.py     |  11 +-
 .../tensorflow/decompositions_test.py         |  11 +-
 6 files changed, 118 insertions(+), 127 deletions(-)

diff --git a/tensornetwork/backends/numpy/decompositions.py b/tensornetwork/backends/numpy/decompositions.py
index 4f47526d0..d45c7275f 100644
--- a/tensornetwork/backends/numpy/decompositions.py
+++ b/tensornetwork/backends/numpy/decompositions.py
@@ -125,25 +125,19 @@ def rq(
 
 
 def cholesky(
-    tf: Any,
+    np: Any,
     tensor: Tensor,
-    pivot_axis: int,
-    non_negative_diagonal: bool
+    pivot_axis: int
 ) -> 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))
+  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
diff --git a/tensornetwork/backends/numpy/decompositions_test.py b/tensornetwork/backends/numpy/decompositions_test.py
index ac04afdcc..92f780b50 100644
--- a/tensornetwork/backends/numpy/decompositions_test.py
+++ b/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):
 
@@ -54,10 +56,10 @@ def test_qr(self):
 
   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)
+    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)
diff --git a/tensornetwork/backends/pytorch/decompositions.py b/tensornetwork/backends/pytorch/decompositions.py
index 535807e2d..ffcc861b2 100644
--- a/tensornetwork/backends/pytorch/decompositions.py
+++ b/tensornetwork/backends/pytorch/decompositions.py
@@ -124,24 +124,18 @@ def svd(
 def cholesky(
     torch: Any,
     tensor: Tensor,
-    pivot_axis: int,
-    non_negative_diagonal: bool = False
+    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:]
-  
+  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 = 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])
+  L = np.linalg.cholesky(tensor)
   return L
 
 def qr(
diff --git a/tensornetwork/backends/pytorch/decompositions_test.py b/tensornetwork/backends/pytorch/decompositions_test.py
index bafcd2447..ecdb87fe0 100644
--- a/tensornetwork/backends/pytorch/decompositions_test.py
+++ b/tensornetwork/backends/pytorch/decompositions_test.py
@@ -16,99 +16,102 @@
 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))
-  u, s, vh, _ = decompositions.svd(torch, val, 2)
-  assert u.shape == (2, 3, 6)
-  assert s.shape == (6,)
-  np.testing.assert_allclose(s, np.zeros(6))
-  assert vh.shape == (6, 4, 5)
+# def test_expected_shapes():
+#   val = torch.zeros((2, 3, 4, 5))
+#   u, s, vh, _ = decompositions.svd(torch, val, 2)
+#   assert u.shape == (2, 3, 6)
+#   assert s.shape == (6,)
+#   np.testing.assert_allclose(s, np.zeros(6))
+#   assert vh.shape == (6, 4, 5)
 
 
-def test_expected_shapes_qr():
-  val = torch.zeros((2, 3, 4, 5))
-  for non_negative_diagonal in [True, False]:
-    q, r = decompositions.qr(torch, val, 2, non_negative_diagonal)
-    assert q.shape == (2, 3, 6)
-    assert r.shape == (6, 4, 5)
+# def test_expected_shapes_qr():
+#   val = torch.zeros((2, 3, 4, 5))
+#   for non_negative_diagonal in [True, False]:
+#     q, r = decompositions.qr(torch, val, 2, non_negative_diagonal)
+#     assert q.shape == (2, 3, 6)
+#     assert r.shape == (6, 4, 5)
 
 
-def test_expected_shapes_rq():
-  val = torch.zeros((2, 3, 4, 5))
-  for non_negative_diagonal in [True, False]:
-    r, q = decompositions.rq(torch, val, 2, non_negative_diagonal)
-    assert r.shape == (2, 3, 6)
-    assert q.shape == (6, 4, 5)
+# def test_expected_shapes_rq():
+#   val = torch.zeros((2, 3, 4, 5))
+#   for non_negative_diagonal in [True, False]:
+#     r, q = decompositions.rq(torch, val, 2, non_negative_diagonal)
+#     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.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)
-  for non_negative_diagonal in [True, False]:
-    r, q = decompositions.rq(torch, random_matrix, 1, non_negative_diagonal)
-    np.testing.assert_allclose(r.mm(q), random_matrix)
-
-
-def test_qr():
-  random_matrix = torch.rand([10, 10], dtype=torch.float64)
-  for non_negative_diagonal in [True, False]:
-    q, r = decompositions.rq(torch, random_matrix, 1, non_negative_diagonal)
-    np.testing.assert_allclose(q.mm(r), random_matrix)
-
-
-def test_max_singular_values():
-  np.random.seed(2018)
   random_matrix = np.random.rand(10, 10)
-  unitary1, _, unitary2 = np.linalg.svd(random_matrix)
-  singular_values = np.array(range(10))
-  val = unitary1.dot(np.diag(singular_values).dot(unitary2.T))
-  u, s, vh, trun = decompositions.svd(
-      torch, torch.tensor(val), 1, max_singular_values=7)
-  assert u.shape == (10, 7)
-  assert s.shape == (7,)
-  np.testing.assert_array_almost_equal(s, np.arange(9, 2, -1))
-  assert vh.shape == (7, 10)
-  np.testing.assert_array_almost_equal(trun, np.arange(2, -1, -1))
-
-
-def test_max_truncation_error():
-  np.random.seed(2019)
-  random_matrix = np.random.rand(10, 10)
-  unitary1, _, unitary2 = np.linalg.svd(random_matrix)
-  singular_values = np.array(range(10))
-  val = unitary1.dot(np.diag(singular_values).dot(unitary2.T))
-  u, s, vh, trun = decompositions.svd(
-      torch, torch.Tensor(val), 1, max_truncation_error=math.sqrt(5.1))
-  assert u.shape == (10, 7)
-  assert s.shape == (7,)
-  np.testing.assert_array_almost_equal(s, np.arange(9, 2, -1), decimal=5)
-  assert vh.shape == (7, 10)
-  np.testing.assert_array_almost_equal(trun, np.arange(2, -1, -1))
-
-
-def test_max_truncation_error_relative():
-  absolute = np.diag([2.0, 1.0, 0.2, 0.1])
-  relative = np.diag([2.0, 1.0, 0.2, 0.1])
-  max_truncation_err = 0.2
-  _, _, _, trunc_sv_absolute = decompositions.svd(
-      torch,
-      torch.Tensor(absolute),
-      1,
-      max_truncation_error=max_truncation_err,
-      relative=False)
-  _, _, _, trunc_sv_relative = decompositions.svd(
-      torch,
-      torch.Tensor(relative),
-      1,
-      max_truncation_error=max_truncation_err,
-      relative=True)
-  np.testing.assert_almost_equal(trunc_sv_absolute, [0.1])
-  np.testing.assert_almost_equal(trunc_sv_relative, [0.2, 0.1])
+  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)
+#   for non_negative_diagonal in [True, False]:
+#     r, q = decompositions.rq(torch, random_matrix, 1, non_negative_diagonal)
+#     np.testing.assert_allclose(r.mm(q), random_matrix)
+
+
+# def test_qr():
+#   random_matrix = torch.rand([10, 10], dtype=torch.float64)
+#   for non_negative_diagonal in [True, False]:
+#     q, r = decompositions.rq(torch, random_matrix, 1, non_negative_diagonal)
+#     np.testing.assert_allclose(q.mm(r), random_matrix)
+
+
+# def test_max_singular_values():
+#   np.random.seed(2018)
+#   random_matrix = np.random.rand(10, 10)
+#   unitary1, _, unitary2 = np.linalg.svd(random_matrix)
+#   singular_values = np.array(range(10))
+#   val = unitary1.dot(np.diag(singular_values).dot(unitary2.T))
+#   u, s, vh, trun = decompositions.svd(
+#       torch, torch.tensor(val), 1, max_singular_values=7)
+#   assert u.shape == (10, 7)
+#   assert s.shape == (7,)
+#   np.testing.assert_array_almost_equal(s, np.arange(9, 2, -1))
+#   assert vh.shape == (7, 10)
+#   np.testing.assert_array_almost_equal(trun, np.arange(2, -1, -1))
+
+
+# def test_max_truncation_error():
+#   np.random.seed(2019)
+#   random_matrix = np.random.rand(10, 10)
+#   unitary1, _, unitary2 = np.linalg.svd(random_matrix)
+#   singular_values = np.array(range(10))
+#   val = unitary1.dot(np.diag(singular_values).dot(unitary2.T))
+#   u, s, vh, trun = decompositions.svd(
+#       torch, torch.Tensor(val), 1, max_truncation_error=math.sqrt(5.1))
+#   assert u.shape == (10, 7)
+#   assert s.shape == (7,)
+#   np.testing.assert_array_almost_equal(s, np.arange(9, 2, -1), decimal=5)
+#   assert vh.shape == (7, 10)
+#   np.testing.assert_array_almost_equal(trun, np.arange(2, -1, -1))
+
+
+# def test_max_truncation_error_relative():
+#   absolute = np.diag([2.0, 1.0, 0.2, 0.1])
+#   relative = np.diag([2.0, 1.0, 0.2, 0.1])
+#   max_truncation_err = 0.2
+#   _, _, _, trunc_sv_absolute = decompositions.svd(
+#       torch,
+#       torch.Tensor(absolute),
+#       1,
+#       max_truncation_error=max_truncation_err,
+#       relative=False)
+#   _, _, _, trunc_sv_relative = decompositions.svd(
+#       torch,
+#       torch.Tensor(relative),
+#       1,
+#       max_truncation_error=max_truncation_err,
+#       relative=True)
+#   np.testing.assert_almost_equal(trunc_sv_absolute, [0.1])
+#   np.testing.assert_almost_equal(trunc_sv_relative, [0.2, 0.1])
diff --git a/tensornetwork/backends/tensorflow/decompositions.py b/tensornetwork/backends/tensorflow/decompositions.py
index 4f92a2bbe..f819455ad 100644
--- a/tensornetwork/backends/tensorflow/decompositions.py
+++ b/tensornetwork/backends/tensorflow/decompositions.py
@@ -231,8 +231,7 @@ def rq(
 def cholesky(
     tf: Any,
     tensor: Tensor,
-    pivot_axis: int,
-    non_negative_diagonal: bool
+    pivot_axis: int
 ) -> Tuple[Tensor, Tensor]:
   """ Computes de cholesky decomposition of a tensor.
 
@@ -245,9 +244,5 @@ def cholesky(
                       [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
\ No newline at end of file
+  return L
+  
\ No newline at end of file
diff --git a/tensornetwork/backends/tensorflow/decompositions_test.py b/tensornetwork/backends/tensorflow/decompositions_test.py
index 87042264d..ee8c03b22 100644
--- a/tensornetwork/backends/tensorflow/decompositions_test.py
+++ b/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,12 +56,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)
 
+  # @pytest.mark.parametrize("dtype", np_dtypes)
   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)
+    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)

From 367ce3d64c68cff203754eecf8322d4a5115e221 Mon Sep 17 00:00:00 2001
From: LuiGiovanni <giovanni.lopez3523@gmail.com>
Date: Mon, 28 Dec 2020 12:39:42 -0700
Subject: [PATCH 21/21] Fixed commented testing functions

---
 .../backends/pytorch/decompositions_test.py   | 170 +++++++++---------
 1 file changed, 85 insertions(+), 85 deletions(-)

diff --git a/tensornetwork/backends/pytorch/decompositions_test.py b/tensornetwork/backends/pytorch/decompositions_test.py
index ecdb87fe0..4d157a4a9 100644
--- a/tensornetwork/backends/pytorch/decompositions_test.py
+++ b/tensornetwork/backends/pytorch/decompositions_test.py
@@ -20,29 +20,29 @@
 
 np_dtypes = [np.float64, np.complex128]
 
-# def test_expected_shapes():
-#   val = torch.zeros((2, 3, 4, 5))
-#   u, s, vh, _ = decompositions.svd(torch, val, 2)
-#   assert u.shape == (2, 3, 6)
-#   assert s.shape == (6,)
-#   np.testing.assert_allclose(s, np.zeros(6))
-#   assert vh.shape == (6, 4, 5)
-
-
-# def test_expected_shapes_qr():
-#   val = torch.zeros((2, 3, 4, 5))
-#   for non_negative_diagonal in [True, False]:
-#     q, r = decompositions.qr(torch, val, 2, non_negative_diagonal)
-#     assert q.shape == (2, 3, 6)
-#     assert r.shape == (6, 4, 5)
-
-
-# def test_expected_shapes_rq():
-#   val = torch.zeros((2, 3, 4, 5))
-#   for non_negative_diagonal in [True, False]:
-#     r, q = decompositions.rq(torch, val, 2, non_negative_diagonal)
-#     assert r.shape == (2, 3, 6)
-#     assert q.shape == (6, 4, 5)
+def test_expected_shapes():
+  val = torch.zeros((2, 3, 4, 5))
+  u, s, vh, _ = decompositions.svd(torch, val, 2)
+  assert u.shape == (2, 3, 6)
+  assert s.shape == (6,)
+  np.testing.assert_allclose(s, np.zeros(6))
+  assert vh.shape == (6, 4, 5)
+
+
+def test_expected_shapes_qr():
+  val = torch.zeros((2, 3, 4, 5))
+  for non_negative_diagonal in [True, False]:
+    q, r = decompositions.qr(torch, val, 2, non_negative_diagonal)
+    assert q.shape == (2, 3, 6)
+    assert r.shape == (6, 4, 5)
+
+
+def test_expected_shapes_rq():
+  val = torch.zeros((2, 3, 4, 5))
+  for non_negative_diagonal in [True, False]:
+    r, q = decompositions.rq(torch, val, 2, non_negative_diagonal)
+    assert r.shape == (2, 3, 6)
+    assert q.shape == (6, 4, 5)
 
 # @pytest.mark.parametrize("dtype", np_dtypes)
 def test_cholesky():
@@ -53,65 +53,65 @@ def test_cholesky():
   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)
-#   for non_negative_diagonal in [True, False]:
-#     r, q = decompositions.rq(torch, random_matrix, 1, non_negative_diagonal)
-#     np.testing.assert_allclose(r.mm(q), random_matrix)
-
-
-# def test_qr():
-#   random_matrix = torch.rand([10, 10], dtype=torch.float64)
-#   for non_negative_diagonal in [True, False]:
-#     q, r = decompositions.rq(torch, random_matrix, 1, non_negative_diagonal)
-#     np.testing.assert_allclose(q.mm(r), random_matrix)
-
-
-# def test_max_singular_values():
-#   np.random.seed(2018)
-#   random_matrix = np.random.rand(10, 10)
-#   unitary1, _, unitary2 = np.linalg.svd(random_matrix)
-#   singular_values = np.array(range(10))
-#   val = unitary1.dot(np.diag(singular_values).dot(unitary2.T))
-#   u, s, vh, trun = decompositions.svd(
-#       torch, torch.tensor(val), 1, max_singular_values=7)
-#   assert u.shape == (10, 7)
-#   assert s.shape == (7,)
-#   np.testing.assert_array_almost_equal(s, np.arange(9, 2, -1))
-#   assert vh.shape == (7, 10)
-#   np.testing.assert_array_almost_equal(trun, np.arange(2, -1, -1))
-
-
-# def test_max_truncation_error():
-#   np.random.seed(2019)
-#   random_matrix = np.random.rand(10, 10)
-#   unitary1, _, unitary2 = np.linalg.svd(random_matrix)
-#   singular_values = np.array(range(10))
-#   val = unitary1.dot(np.diag(singular_values).dot(unitary2.T))
-#   u, s, vh, trun = decompositions.svd(
-#       torch, torch.Tensor(val), 1, max_truncation_error=math.sqrt(5.1))
-#   assert u.shape == (10, 7)
-#   assert s.shape == (7,)
-#   np.testing.assert_array_almost_equal(s, np.arange(9, 2, -1), decimal=5)
-#   assert vh.shape == (7, 10)
-#   np.testing.assert_array_almost_equal(trun, np.arange(2, -1, -1))
-
-
-# def test_max_truncation_error_relative():
-#   absolute = np.diag([2.0, 1.0, 0.2, 0.1])
-#   relative = np.diag([2.0, 1.0, 0.2, 0.1])
-#   max_truncation_err = 0.2
-#   _, _, _, trunc_sv_absolute = decompositions.svd(
-#       torch,
-#       torch.Tensor(absolute),
-#       1,
-#       max_truncation_error=max_truncation_err,
-#       relative=False)
-#   _, _, _, trunc_sv_relative = decompositions.svd(
-#       torch,
-#       torch.Tensor(relative),
-#       1,
-#       max_truncation_error=max_truncation_err,
-#       relative=True)
-#   np.testing.assert_almost_equal(trunc_sv_absolute, [0.1])
-#   np.testing.assert_almost_equal(trunc_sv_relative, [0.2, 0.1])
+def test_rq():
+  random_matrix = torch.rand([10, 10], dtype=torch.float64)
+  for non_negative_diagonal in [True, False]:
+    r, q = decompositions.rq(torch, random_matrix, 1, non_negative_diagonal)
+    np.testing.assert_allclose(r.mm(q), random_matrix)
+
+
+def test_qr():
+  random_matrix = torch.rand([10, 10], dtype=torch.float64)
+  for non_negative_diagonal in [True, False]:
+    q, r = decompositions.rq(torch, random_matrix, 1, non_negative_diagonal)
+    np.testing.assert_allclose(q.mm(r), random_matrix)
+
+
+def test_max_singular_values():
+  np.random.seed(2018)
+  random_matrix = np.random.rand(10, 10)
+  unitary1, _, unitary2 = np.linalg.svd(random_matrix)
+  singular_values = np.array(range(10))
+  val = unitary1.dot(np.diag(singular_values).dot(unitary2.T))
+  u, s, vh, trun = decompositions.svd(
+      torch, torch.tensor(val), 1, max_singular_values=7)
+  assert u.shape == (10, 7)
+  assert s.shape == (7,)
+  np.testing.assert_array_almost_equal(s, np.arange(9, 2, -1))
+  assert vh.shape == (7, 10)
+  np.testing.assert_array_almost_equal(trun, np.arange(2, -1, -1))
+
+
+def test_max_truncation_error():
+  np.random.seed(2019)
+  random_matrix = np.random.rand(10, 10)
+  unitary1, _, unitary2 = np.linalg.svd(random_matrix)
+  singular_values = np.array(range(10))
+  val = unitary1.dot(np.diag(singular_values).dot(unitary2.T))
+  u, s, vh, trun = decompositions.svd(
+      torch, torch.Tensor(val), 1, max_truncation_error=math.sqrt(5.1))
+  assert u.shape == (10, 7)
+  assert s.shape == (7,)
+  np.testing.assert_array_almost_equal(s, np.arange(9, 2, -1), decimal=5)
+  assert vh.shape == (7, 10)
+  np.testing.assert_array_almost_equal(trun, np.arange(2, -1, -1))
+
+
+def test_max_truncation_error_relative():
+  absolute = np.diag([2.0, 1.0, 0.2, 0.1])
+  relative = np.diag([2.0, 1.0, 0.2, 0.1])
+  max_truncation_err = 0.2
+  _, _, _, trunc_sv_absolute = decompositions.svd(
+      torch,
+      torch.Tensor(absolute),
+      1,
+      max_truncation_error=max_truncation_err,
+      relative=False)
+  _, _, _, trunc_sv_relative = decompositions.svd(
+      torch,
+      torch.Tensor(relative),
+      1,
+      max_truncation_error=max_truncation_err,
+      relative=True)
+  np.testing.assert_almost_equal(trunc_sv_absolute, [0.1])
+  np.testing.assert_almost_equal(trunc_sv_relative, [0.2, 0.1])