add broadcasting support to backends (#513)

* adding broadcast_right_multiplication to backends

* fix test

Author: mganahl

tensornetwork/backends/base_backend.py

@@ -354,16 +354,16 @@ def eigs(self,
     raise NotImplementedError("Backend '{}' has not implemented eigs.".format(
         self.name))
 
-  def eigsh_lanczos(self,
-                    A: Callable,
-                    initial_state: Optional[Tensor] = None,
-                    num_krylov_vecs: Optional[int] = 200,
-                    numeig: Optional[int] = 1,
-                    tol: Optional[float] = 1E-8,
-                    delta: Optional[float] = 1E-8,
-                    ndiag: Optional[int] = 20,
-                    reorthogonalize: Optional[bool] = False
-                   ) -> Tuple[List, List]:
+  def eigsh_lanczos(
+      self,
+      A: Callable,
+      initial_state: Optional[Tensor] = None,
+      num_krylov_vecs: Optional[int] = 200,
+      numeig: Optional[int] = 1,
+      tol: Optional[float] = 1E-8,
+      delta: Optional[float] = 1E-8,
+      ndiag: Optional[int] = 20,
+      reorthogonalize: Optional[bool] = False) -> Tuple[List, List]:
     """
     Lanczos method for finding the lowest eigenvector-eigenvalue pairs
     of `A`. 
@@ -444,8 +444,8 @@ def divide(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
       Returns:
         Tensor
     """
-    raise NotImplementedError(
-        "Backend '{}' has not implemented divide.".format(self.name))
+    raise NotImplementedError("Backend '{}' has not implemented divide.".format(
+        self.name))
 
   def index_update(self, tensor: Tensor, mask: Tensor,
                    assignee: Tensor) -> Tensor:
@@ -471,3 +471,17 @@ def inv(self, matrix: Tensor) -> Tensor:
     """
     raise NotImplementedError("Backend '{}' has not implemented `inv`.".format(
         self.name))
+
+  def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
+    """
+    Perform broadcasting for multiplication of `tensor2` onto `tensor1`, i.e.
+    `tensor1` * tensor2`.
+    Args:
+      tensor1: A tensor.
+      tensor2: A tensor.
+    Returns:
+      Tensor: The result of multiplying `tensor1` onto `tensor2`.
+    """
+    raise NotImplementedError(
+        "Backend '{}' has not implemented `broadcast_right_multiplication`."
+        .format(self.name))

tensornetwork/backends/jax/jax_backend_test.py

@@ -241,8 +241,8 @@ def test_random_uniform_boundaries(dtype):
   backend = jax_backend.JaxBackend()
   a = backend.random_uniform((4, 4), seed=10, dtype=dtype)
   b = backend.random_uniform((4, 4), (lb, ub), seed=10, dtype=dtype)
-  assert((a >= 0).all() and (a <= 1).all() and
-         (b >= lb).all() and (b <= ub).all())
+  assert ((a >= 0).all() and (a <= 1).all() and (b >= lb).all() and
+          (b <= ub).all())
 
 
 def test_random_uniform_behavior():
@@ -296,3 +296,20 @@ def test_index_update(dtype):
   np_tensor = np.array(tensor)
   np_tensor[np_tensor > 0.1] = 0.0
   np.testing.assert_allclose(out, np_tensor)
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
+def test_broadcast_right_multiplication(dtype):
+  backend = jax_backend.JaxBackend()
+  tensor1 = backend.randn((2, 3), dtype=dtype, seed=10)
+  tensor2 = backend.randn((3,), dtype=dtype, seed=10)
+  out = backend.broadcast_right_multiplication(tensor1, tensor2)
+  np.testing.assert_allclose(out, np.array(tensor1) * np.array(tensor2))
+
+
+def test_broadcast_right_multiplication_raises():
+  backend = jax_backend.JaxBackend()
+  tensor1 = backend.randn((2, 3))
+  tensor2 = backend.randn((3, 3))
+  with pytest.raises(ValueError):
+    backend.broadcast_right_multiplication(tensor1, tensor2)

tensornetwork/backends/numpy/numpy_backend.py

@@ -44,10 +44,13 @@ def svd_decomposition(self,
                         max_truncation_error: Optional[float] = None,
                         relative: Optional[bool] = False
                        ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
-    return decompositions.svd_decomposition(self.np, tensor, split_axis,
-                                            max_singular_values,
-                                            max_truncation_error,
-                                            relative=relative)
+    return decompositions.svd_decomposition(
+        self.np,
+        tensor,
+        split_axis,
+        max_singular_values,
+        max_truncation_error,
+        relative=relative)
 
   def qr_decomposition(
       self,
@@ -145,12 +148,12 @@ def random_uniform(self,
     dtype = dtype if dtype is not None else self.np.float64
     if ((self.np.dtype(dtype) is self.np.dtype(self.np.complex128)) or
         (self.np.dtype(dtype) is self.np.dtype(self.np.complex64))):
-      return self.np.random.uniform(boundaries[0], boundaries[1], shape).astype(
-          dtype) + 1j * self.np.random.uniform(boundaries[0],
-                                               boundaries[1],
-                                               shape).astype(dtype)
-    return self.np.random.uniform(boundaries[0],
-                                  boundaries[1], shape).astype(dtype)
+      return self.np.random.uniform(
+          boundaries[0], boundaries[1],
+          shape).astype(dtype) + 1j * self.np.random.uniform(
+              boundaries[0], boundaries[1], shape).astype(dtype)
+    return self.np.random.uniform(boundaries[0], boundaries[1],
+                                  shape).astype(dtype)
 
   def conj(self, tensor: Tensor) -> Tensor:
     return self.np.conj(tensor)
@@ -244,16 +247,16 @@ def eigs(self,
       U = U.astype(dtype)
     return list(eta), [U[:, n] for n in range(numeig)]
 
-  def eigsh_lanczos(self,
-                    A: Callable,
-                    initial_state: Optional[Tensor] = None,
-                    num_krylov_vecs: Optional[int] = 200,
-                    numeig: Optional[int] = 1,
-                    tol: Optional[float] = 1E-8,
-                    delta: Optional[float] = 1E-8,
-                    ndiag: Optional[int] = 20,
-                    reorthogonalize: Optional[bool] = False
-                   ) -> Tuple[List, List]:
+  def eigsh_lanczos(
+      self,
+      A: Callable,
+      initial_state: Optional[Tensor] = None,
+      num_krylov_vecs: Optional[int] = 200,
+      numeig: Optional[int] = 1,
+      tol: Optional[float] = 1E-8,
+      delta: Optional[float] = 1E-8,
+      ndiag: Optional[int] = 20,
+      reorthogonalize: Optional[bool] = False) -> Tuple[List, List]:
     """
     Lanczos method for finding the lowest eigenvector-eigenvalue pairs
     of a linear operator `A`. If no `initial_state` is provided
@@ -396,3 +399,9 @@ def inv(self, matrix: Tensor) -> Tensor:
       raise ValueError("input to numpy backend method `inv` has shape {}."
                        " Only matrices are supported.".format(matrix.shape))
     return self.np.linalg.inv(matrix)
+
+  def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
+    if len(tensor2.shape) != 1:
+      raise ValueError("only order-1 tensors are allowed for `tensor2`,"
+                       " found `tensor2.shape = {}`".format(tensor2.shape))
+    return tensor1 * tensor2

tensornetwork/backends/numpy/numpy_backend_test.py

@@ -239,8 +239,8 @@ def test_random_uniform_boundaries(dtype):
   backend = numpy_backend.NumPyBackend()
   a = backend.random_uniform((4, 4), seed=10, dtype=dtype)
   b = backend.random_uniform((4, 4), (lb, ub), seed=10, dtype=dtype)
-  assert((a >= 0).all() and (a <= 1).all() and
-         (b >= lb).all() and (b <= ub).all())
+  assert ((a >= 0).all() and (a <= 1).all() and (b >= lb).all() and
+          (b <= ub).all())
 
 
 def test_random_uniform_behavior():
@@ -329,8 +329,8 @@ def __call__(self, x):
       return np.dot(H, x)
 
   mv = LinearOperator(shape=((D,), (D,)), dtype=dtype)
-  eta1, U1 = backend.eigsh_lanczos(mv, reorthogonalize=True, ndiag=1,
-                                   tol=10**(-12), delta=10**(-12))
+  eta1, U1 = backend.eigsh_lanczos(
+      mv, reorthogonalize=True, ndiag=1, tol=10**(-12), delta=10**(-12))
   eta2, U2 = np.linalg.eigh(H)
   v2 = U2[:, 0]
   v2 = v2 / sum(v2)
@@ -353,7 +353,7 @@ def test_eigsh_lanczos_raises():
 def test_eigsh_lanczos_raises_error_for_incompatible_shapes():
   backend = numpy_backend.NumPyBackend()
   A = backend.randn((4, 4), dtype=np.float64)
-  init = backend.randn((3, ), dtype=np.float64)
+  init = backend.randn((3,), dtype=np.float64)
   with pytest.raises(ValueError):
     backend.eigsh_lanczos(A, initial_state=init)
 
@@ -371,7 +371,7 @@ def test_eigsh_lanczos_raises_error_for_untyped_A():
 def test_eigsh_lanczos_raises_error_for_bad_initial_state():
   backend = numpy_backend.NumPyBackend()
   D = 16
-  init = [1]*D
+  init = [1] * D
   M = backend.randn((D, D), dtype=np.float64)
 
   def mv(x):
@@ -383,9 +383,10 @@ def mv(x):
 
 @pytest.mark.parametrize("a, b, expected", [
     pytest.param(1, 1, 2),
-    pytest.param(1., np.ones((1, 2, 3)), 2*np.ones((1, 2, 3))),
-    pytest.param(2.*np.ones(()), 1., 3.*np.ones((1, 2, 3))),
-    pytest.param(2.*np.ones(()), 1.*np.ones((1, 2, 3)), 3.*np.ones((1, 2, 3))),
+    pytest.param(1., np.ones((1, 2, 3)), 2 * np.ones((1, 2, 3))),
+    pytest.param(2. * np.ones(()), 1., 3. * np.ones((1, 2, 3))),
+    pytest.param(2. * np.ones(()), 1. * np.ones((1, 2, 3)), 3. * np.ones(
+        (1, 2, 3))),
 ])
 def test_addition(a, b, expected):
   backend = numpy_backend.NumPyBackend()
@@ -399,7 +400,7 @@ def test_addition(a, b, expected):
 
 @pytest.mark.parametrize("a, b, expected", [
     pytest.param(1, 1, 0),
-    pytest.param(2., 1.*np.ones((1, 2, 3)), 1.*np.ones((1, 2, 3))),
+    pytest.param(2., 1. * np.ones((1, 2, 3)), 1. * np.ones((1, 2, 3))),
     pytest.param(np.ones((1, 2, 3)), 1., np.zeros((1, 2, 3))),
     pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), np.zeros((1, 2, 3))),
 ])
@@ -415,7 +416,7 @@ def test_subtraction(a, b, expected):
 
 @pytest.mark.parametrize("a, b, expected", [
     pytest.param(1, 1, 1),
-    pytest.param(2., 1.*np.ones((1, 2, 3)), 2.*np.ones((1, 2, 3))),
+    pytest.param(2., 1. * np.ones((1, 2, 3)), 2. * np.ones((1, 2, 3))),
     pytest.param(np.ones((1, 2, 3)), 1., np.ones((1, 2, 3))),
     pytest.param(np.ones((1, 2, 3)), np.ones((1, 2, 3)), np.ones((1, 2, 3))),
 ])
@@ -431,9 +432,10 @@ def test_multiply(a, b, expected):
 
 @pytest.mark.parametrize("a, b, expected", [
     pytest.param(2., 2., 1.),
-    pytest.param(2., 0.5*np.ones((1, 2, 3)), 4.*np.ones((1, 2, 3))),
-    pytest.param(np.ones(()), 2., 0.5*np.ones((1, 2, 3))),
-    pytest.param(np.ones(()), 2.*np.ones((1, 2, 3)), 0.5*np.ones((1, 2, 3))),
+    pytest.param(2., 0.5 * np.ones((1, 2, 3)), 4. * np.ones((1, 2, 3))),
+    pytest.param(np.ones(()), 2., 0.5 * np.ones((1, 2, 3))),
+    pytest.param(
+        np.ones(()), 2. * np.ones((1, 2, 3)), 0.5 * np.ones((1, 2, 3))),
 ])
 def test_divide(a, b, expected):
   backend = numpy_backend.NumPyBackend()
@@ -531,7 +533,7 @@ def test_eigs_raises_error_for_unsupported_which(which):
 def test_eigs_raises_error_for_incompatible_shapes():
   backend = numpy_backend.NumPyBackend()
   A = backend.randn((4, 4), dtype=np.float64)
-  init = backend.randn((3, ), dtype=np.float64)
+  init = backend.randn((3,), dtype=np.float64)
   with pytest.raises(ValueError):
     backend.eigs(A, initial_state=init)
 
@@ -559,7 +561,7 @@ def test_eigs_raises_error_for_untyped_A():
 def test_eigs_raises_error_for_bad_initial_state():
   backend = numpy_backend.NumPyBackend()
   D = 16
-  init = [1]*D
+  init = [1] * D
   M = backend.randn((D, D), dtype=np.float64)
 
   def mv(x):
@@ -609,3 +611,21 @@ def test_matrix_inv_raises(dtype):
   matrix = backend.randn((4, 4, 4), dtype=dtype, seed=10)
   with pytest.raises(ValueError):
     backend.inv(matrix)
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
+def test_broadcast_right_multiplication(dtype):
+  backend = numpy_backend.NumPyBackend()
+  tensor1 = backend.randn((2, 4, 3), dtype=dtype, seed=10)
+  tensor2 = backend.randn((3,), dtype=dtype, seed=10)
+  out = backend.broadcast_right_multiplication(tensor1, tensor2)
+  np.testing.assert_allclose(out, tensor1 * tensor2)
+
+
+def test_broadcast_right_multiplication_raises():
+  dtype = np.float64
+  backend = numpy_backend.NumPyBackend()
+  tensor1 = backend.randn((2, 4, 3), dtype=dtype, seed=10)
+  tensor2 = backend.randn((3, 3), dtype=dtype, seed=10)
+  with pytest.raises(ValueError):
+    backend.broadcast_right_multiplication(tensor1, tensor2)

tensornetwork/backends/pytorch/pytorch_backend.py

@@ -52,10 +52,13 @@ def svd_decomposition(self,
                         max_truncation_error: Optional[float] = None,
                         relative: Optional[bool] = False
                        ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
-    return decompositions.svd_decomposition(self.torch, tensor, split_axis,
-                                            max_singular_values,
-                                            max_truncation_error,
-                                            relative=relative)
+    return decompositions.svd_decomposition(
+        self.torch,
+        tensor,
+        split_axis,
+        max_singular_values,
+        max_truncation_error,
+        relative=relative)
 
   def qr_decomposition(
       self,
@@ -158,16 +161,16 @@ def eigs(self,
     raise NotImplementedError("Backend '{}' has not implemented eigs.".format(
         self.name))
 
-  def eigsh_lanczos(self,
-                    A: Callable,
-                    initial_state: Optional[Tensor] = None,
-                    num_krylov_vecs: Optional[int] = 200,
-                    numeig: Optional[int] = 1,
-                    tol: Optional[float] = 1E-8,
-                    delta: Optional[float] = 1E-8,
-                    ndiag: Optional[int] = 20,
-                    reorthogonalize: Optional[bool] = False
-                   ) -> Tuple[List, List]:
+  def eigsh_lanczos(
+      self,
+      A: Callable,
+      initial_state: Optional[Tensor] = None,
+      num_krylov_vecs: Optional[int] = 200,
+      numeig: Optional[int] = 1,
+      tol: Optional[float] = 1E-8,
+      delta: Optional[float] = 1E-8,
+      ndiag: Optional[int] = 20,
+      reorthogonalize: Optional[bool] = False) -> Tuple[List, List]:
     """
     Lanczos method for finding the lowest eigenvector-eigenvalue pairs
     of a `LinearOperator` `A`.
@@ -302,3 +305,11 @@ def inv(self, matrix: Tensor) -> Tensor:
           "input to pytorch backend method `inv` has shape {}. Only matrices are supported."
           .format(matrix.shape))
     return matrix.inverse()
+
+  def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
+    if len(tensor2.shape) != 1:
+      raise ValueError(
+          "only order-1 tensors are allowed for `tensor2`, found `tensor2.shape = {}`"
+          .format(tensor2.shape))
+
+    return tensor1 * tensor2