Added Cholskey NotImplementedError function and test

tensornetwork/backends/abstract_backend.py

@@ -144,6 +144,13 @@ def qr(self,
     raise NotImplementedError("Backend '{}' has not implemented qr.".format(
         self.name))
 
+  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 rq(self,
          tensor: Tensor,
          pivot_axis: int = -1,

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):

tensornetwork/backends/jax/jax_backend.py

@@ -874,3 +874,17 @@ def sign(self, tensor: Tensor) -> Tensor:
 
   def item(self, tensor):
     return tensor.item()
+
+  def power(self, a: Tensor, b: Union[Tensor, int]) -> 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)

tensornetwork/backends/jax/jax_backend_test.py

@@ -974,7 +974,6 @@ def matvec_jax(vector, matrix):
       num_krylov_vecs=100,
       tol=0.0001)
 
-
 def test_sum():
   np.random.seed(10)
   backend = jax_backend.JaxBackend()
@@ -1240,3 +1239,18 @@ 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 = 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 = jax.numpy.power(base_tensor, power)
+  np.testing.assert_allclose(expected, actual)