Support of .imag and .real for backends

tensornetwork/backends/abstract_backend.py

@@ -1010,4 +1010,25 @@ def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
     """
     raise NotImplementedError(
         f"Backend {self.name} has not implemented power.")
-
+
+  def real(self, tensor: Tensor) -> Tensor:
+      """Return Re(tensor)
+      Args:
+          tensor: Input Tensor
+
+      Returns: Re(tensor), real part of tensor
+
+      """
+      raise NotImplementedError(
+          f"Backend {self.name} has not implemented .real")
+
+  def imag(self, tensor: Tensor) -> Tensor:
+      """Return Im(tensor)
+      Args:
+          tensor: Input tensor.
+
+      Returns: Im(tensor), imaginary part of tensor.
+
+      """
+      raise NotImplementedError(
+          f"Backend {self.name} has not implemented .imag")

tensornetwork/backends/jax/jax_backend.py

@@ -871,3 +871,25 @@ def sign(self, tensor: Tensor) -> Tensor:
       tensor: The input tensor.
     """
     return jnp.sign(tensor)
+
+  def real(self, tensor: Tensor) -> Tensor:
+    """
+
+    Args:
+      tensor: Input Tensor
+
+    Returns: Re(tensor),Real part of tensor
+
+    """
+    return jnp.real(tensor)
+
+  def imag(self, tensor: Tensor) -> Tensor:
+    """
+
+    Args:
+      tensor: Input Tensor
+
+    Returns: Im(tensor), Imaginary part of tensor
+
+    """
+    return jnp.imag(tensor)

tensornetwork/backends/jax/jax_backend_test.py

@@ -1096,3 +1096,27 @@ def test_inv(dtype, atol):
   tensor = backend.randn((10, 10, 10), dtype=dtype, seed=10)
   with pytest.raises(ValueError, match="input to"):
     backend.inv(tensor)
+
+
+@pytest.mark.parametrize("dtype", np_dtypes)
+@pytest.mark.parametrize("input_cur", [np.array([1, 2]),
+                                       np.array([5+6j, 4+10j]),
+                                       np.array([5j, 7j])])
+def test_real(dtype,input_cur):
+  backend = jax_backend.JaxBackend()
+  cur = backend.convert_to_tensor(input_cur)
+  acual = backend.real(cur)
+  expcted = jax.numpy.real(cur)
+  np.testing.assert_allclose(acual, expcted)
+
+
+@pytest.mark.parametrize("dtype", np_dtypes)
+@pytest.mark.parametrize("input_cur", [np.array([1, 2]),
+                                       np.array([5+6j, 4+10j]),
+                                       np.array([5j, 7j])])
+def test_imag(dtype, input_cur):
+  backend = jax_backend.JaxBackend()
+  cur = backend.convert_to_tensor(input_cur)
+  acual = backend.imag(cur)
+  expcted = jax.numpy.imag(cur)
+  np.testing.assert_allclose(acual, expcted)

tensornetwork/backends/numpy/numpy_backend.py

@@ -759,24 +759,25 @@ def deserialize_tensor(self, s: str) -> Tensor:
     m.write(s.encode('latin-1'))
     m.seek(0)
     return np.load(m)
-  
-  def power(self, a: Tensor, b: Union[Tensor, float]) -> Tensor:
+
+  def real(self, tensor: Tensor) -> Tensor:
     """
-    Returns the exponentiation of tensor a raised to b.  
-      If b is a tensor, then the exponentiation is element-wise 
-        between the two tensors, with a as the base and b as the power.
-        Note that a and b must be broadcastable to the same shape if 
-        b is a tensor.
-      If b is a scalar, then the exponentiation is each value in a
-        raised to the power of b.
-
+    Retuns Re(tensor), returns real element in tensor given
+    Args:
+      tensor: Input Tensor
+
+    Returns:
+      returns real element in tensor given
+    """
+    return np.real(tensor)
+
+  def imag(self, tensor: Tensor) -> Tensor:
+    """
+    Returns Im(tensor), returns the Imaginary part of tensor
     Args:
-      a: The tensor containing the bases.
-      b: The tensor containing the powers; or a single scalar as the power.
+      tensor: Input Tensor
 
     Returns:
-      The tensor that is each element of a raised to the 
-        power of b.  Note that the shape of the returned tensor
-        is that produced by the broadcast of a and b.
+      returns the Imaginary part of tensor
     """
-    return np.power(a, b)
+    return np.imag(tensor)

tensornetwork/backends/numpy/numpy_backend_test.py

@@ -950,16 +950,22 @@ def test_serialize(dtype):
   assert isinstance(s, str)
   assert (tensor == backend.deserialize_tensor(s)).all()
 
+
 @pytest.mark.parametrize('dtype', np_dtypes)
-def test_power(dtype):
-  shape = (4, 3, 2)
+@pytest.mark.parametrize("input_cur", [(np.array([1, 2, 3, 4])), (np.array([1j, 2j, 3j, 4j])),
+                                   (np.array([10+2j, 20+2j, 30+5j, 40+20j]))])
+def test_real(dtype, input_cur):
   backend = numpy_backend.NumPyBackend()
-  base_tensor = np.abs(backend.randn(shape, dtype=dtype, seed=10))
-  power_tensor = backend.randn(shape, dtype=dtype, seed=10)
-  actual = backend.power(base_tensor, power_tensor)
-  expected = np.power(base_tensor, power_tensor)
-  np.testing.assert_allclose(expected, actual)
-  power = np.random.rand(1)[0]
-  actual = backend.power(base_tensor, power)
-  expected = np.power(base_tensor, power)
-  np.testing.assert_allclose(expected, actual)
+  cur = backend.convert_to_tensor(input_cur)
+  np.testing.assert_allclose(cur.real, np.real(input_cur))
+
+
+@pytest.mark.parametrize('dtype', np_dtypes)
+@pytest.mark.parametrize("input_cur", [(np.array([1, 2, 3, 4])), (np.array([1j, 2j, 3j, 4j])),
+                         (np.array([10+2j, 20+2j, 30+5j, 40+20j]))])
+def test_imag(dtype, input_cur):
+  backend = numpy_backend.NumPyBackend()
+  cur = backend.convert_to_tensor(input_cur)
+  acual = cur.imag
+  expected = np.imag(input_cur)
+  np.testing.assert_allclose(acual, expected)

tensornetwork/backends/pytorch/pytorch_backend.py

@@ -472,3 +472,35 @@ def sign(self, tensor: Tensor) -> Tensor:
       tensor: The input tensor.
     """
     return torchlib.sign(tensor)
+
+  def real(self, tensor: Tensor) -> Tensor:
+    """Return real part of tensor
+    Args:
+      tensor: Input Tensor
+
+    Returns: Re(tensor),Real part of tensor
+
+    """
+    if torchlib.is_complex(tensor):
+      return torchlib.real(tensor)
+    else:
+      temp = torchlib.zeros_like(tensor, dtype=torchlib.cfloat)
+      tensor = torchlib.as_tensor(tensor, dtype=torchlib.cfloat)
+      torchlib.add(tensor, temp, out=tensor)
+      return torchlib.real(tensor)
+
+  def imag(self,tensor: Tensor) -> Tensor:
+    """Return imag part of tensor
+    Args:
+      tensor: Input tensor
+
+    Returns: Im(tensor),returns Imaginary part of tensor
+
+    """
+    if torchlib.is_complex(tensor):
+      return torchlib.imag(tensor)
+    else:
+      temp = torchlib.zeros_like(tensor, dtype=torchlib.cfloat)
+      tensor = torchlib.as_tensor(tensor, dtype=torchlib.cfloat)
+      torchlib.add(tensor, temp, out=tensor)
+      return torchlib.imag(tensor)

tensornetwork/backends/pytorch/pytorch_backend_test.py

@@ -664,3 +664,29 @@ def test_matmul_rank2():
   actual = backend.matmul(a, b)
   expected = np.matmul(t1, t2)
   np.testing.assert_allclose(expected, actual)
+
+
+@pytest.mark.parametrize("dtype", torch_randn_dtypes)
+@pytest.mark.parametrize("input_cur", [(np.array([1+6j, 2+7j])), (np.array([1j, 2j])),
+                         (np.array([5+3j, 4+7j]))])
+def test_real(dtype, input_cur):
+  backend = pytorch_backend.PyTorchBackend()
+  cur = backend.convert_to_tensor(input_cur)
+  acual = backend.real(cur)
+  expected = torch.real(cur)
+  np.testing.assert_allclose(acual, expected)
+  cur = backend.convert_to_tensor(np.array([1, 2]))
+  np.testing.assert_allclose(backend.real(cur), np.array([1, 2]))
+
+
+@pytest.mark.parametrize("dtype", torch_randn_dtypes)
+@pytest.mark.parametrize("input_cur", [(np.array([1+6j, 2+7j])), (np.array([1j, 2j])),
+                         (np.array([5+3j, 4+7j]))])
+def test_imag(dtype, input_cur):
+  backend = pytorch_backend.PyTorchBackend()
+  cur = backend.convert_to_tensor(input_cur)
+  acual = backend.imag(cur)
+  expected = torch.imag(cur)
+  np.testing.assert_allclose(acual, expected)
+  cur = backend.convert_to_tensor(np.array([1, 2]))
+  np.testing.assert_allclose(backend.imag(cur), np.array([0, 0]))

tensornetwork/backends/symmetric/symmetric_backend.py

@@ -680,3 +680,12 @@ def sign(self, tensor: Tensor) -> Tensor:
 
   def pivot(self, tensor: Tensor, pivot_axis: int = -1) -> Tensor:
     raise NotImplementedError("Symmetric backend doesn't support pivot.")
+
+  def real(self, tensor: Tensor) -> Tensor:
+     temp = tensor.data[abs(tensor.data.imag) > 0]
+     return temp.real
+
+
+  def imag(self, tensor:Tensor) -> Tensor:
+    temp = tensor.data[abs(tensor.data.imag) < 0]
+    return temp.imag

tensornetwork/backends/symmetric/symmetric_backend_test.py

@@ -1644,3 +1644,25 @@ def test_gmres_raises():
     backend.gmres(lambda x: x, b, x0=mps, tol=-0.001)
   with pytest.raises(ValueError, match="atol = "):
     backend.gmres(lambda x: x, b, x0=mps, atol=-0.001)
+
+
+@pytest.mark.parametrize("dtype", np_dtypes)
+@pytest.mark.parametrize("R", [2, 3, 4, 5])
+@pytest.mark.parametrize("num_charges", [1, 2])
+def test_real(dtype, R, num_charges):
+    backend = symmetric_backend.SymmetricBackend()
+    a = get_tensor(R, num_charges, dtype)
+    a = backend.real(a)
+    temp = np.isreal(a)
+    assert(temp.all())
+
+
+@pytest.mark.parametrize("dtype", np_dtypes)
+@pytest.mark.parametrize("R", [2, 3, 4, 5])
+@pytest.mark.parametrize("num_charges", [1, 2])
+def test_imag(dtype, R, num_charges):
+    backend = symmetric_backend.SymmetricBackend()
+    a = get_tensor(R, num_charges, dtype)
+    a = backend.imag(a)
+    temp = np.iscomplex(a)
+    assert(temp.all())

tensornetwork/backends/tensorflow/tensorflow_backend.py

@@ -383,3 +383,25 @@ def sign(self, tensor: Tensor) -> Tensor:
       tensor: The input tensor.
     """
     return tf.math.sign(tensor)
+
+  def real(self, tensor: Tensor) -> Tensor:
+    """
+
+    Args:
+      tensor: Input Tensor
+
+    Returns: Real part of tensor, Re(tensor)
+
+    """
+    return tf.math.real(tensor)
+
+  def imag(self, tensor: Tensor) -> Tensor:
+    """
+
+    Args:
+      tensor: Input Tensor
+
+    Returns: Imaginary part of tensor, Im(tensor)
+
+    """
+    return tf.math.imag(tensor)

tensornetwork/backends/tensorflow/tensorflow_backend_test.py

@@ -599,3 +599,22 @@ def test_pivot(dtype, pivot_axis):
   expected = tf.reshape(tensor, pivot_shape)
   actual = backend.pivot(tensor, pivot_axis=pivot_axis)
   np.testing.assert_allclose(expected, actual)
+
+
+@pytest.mark.parametrize("dtype", tf_dtypes)
+@pytest.mark.parametrize("input_cur", [(np.array([1, 2, 3, 4])), (np.array([1j, 2j, 3j, 4j])),
+                                   (np.array([10+2j, 20+2j, 30+5j, 40+20j]))])
+def test_real(dtype, input_cur):
+  backend = tensorflow_backend.TensorFlowBackend()
+  cur = backend.convert_to_tensor(input_cur)
+  expected = tf.math.real(input_cur)
+  np.testing.assert_allclose(backend.real(cur), expected)
+
+@pytest.mark.parametrize("dtype", tf_dtypes)
+@pytest.mark.parametrize("input_cur", [(np.array([1, 2, 3, 4])), (np.array([1j, 2j, 3j, 4j])),
+                                   (np.array([10+2j, 20+2j, 30+5j, 40+20j]))])
+def test_imag(dtype, input_cur):
+  backend = tensorflow_backend.TensorFlowBackend()
+  cur = backend.convert_to_tensor(input_cur)
+  expected = tf.math.imag(input_cur)
+  np.testing.assert_allclose(backend.imag(cur), expected)