diff --git a/tensornetwork/backends/numpy/numpy_backend.py b/tensornetwork/backends/numpy/numpy_backend.py index b7b8049f9..81f698e13 100644 --- a/tensornetwork/backends/numpy/numpy_backend.py +++ b/tensornetwork/backends/numpy/numpy_backend.py @@ -759,3 +759,24 @@ 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: + """ + 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. + + Args: + a: The tensor containing the bases. + b: The tensor containing the powers; or a single scalar as the power. + + 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. + """ + return np.power(a, b) diff --git a/tensornetwork/backends/numpy/numpy_backend_test.py b/tensornetwork/backends/numpy/numpy_backend_test.py index 59782e5d8..977d01bc6 100644 --- a/tensornetwork/backends/numpy/numpy_backend_test.py +++ b/tensornetwork/backends/numpy/numpy_backend_test.py @@ -949,3 +949,17 @@ def test_serialize(dtype): s = backend.serialize_tensor(tensor) 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) + 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)