Backend broadcast 2 (#514)

* adding broadcast_right_multiplication to backends

* fix test

* bradcasting support added

Author: mganahl

tensornetwork/backends/base_backend.py

@@ -475,7 +475,9 @@ def inv(self, matrix: Tensor) -> Tensor:
   def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
     """
     Perform broadcasting for multiplication of `tensor2` onto `tensor1`, i.e.
-    `tensor1` * tensor2`.
+    `tensor1` * tensor2`, where `tensor1` is an arbitrary tensor and `tensor2` is a
+    one-dimensional tensor. The broadcasting is applied to the last index of 
+    `tensor1`.
     Args:
       tensor1: A tensor.
       tensor2: A tensor.
@@ -485,3 +487,19 @@ def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
     raise NotImplementedError(
         "Backend '{}' has not implemented `broadcast_right_multiplication`."
         .format(self.name))
+
+  def broadcast_left_multiplication(self, tensor1: Tensor, tensor2: Tensor):
+    """
+    Perform broadcasting for multiplication of `tensor1` onto `tensor2`, i.e.
+    `tensor1` * tensor2`, where `tensor2` is an arbitrary tensor and `tensor1` is a
+    one-dimensional tensor. The broadcasting is applied to the first index of 
+    `tensor2`.
+    Args:
+      tensor1: A tensor.
+      tensor2: A tensor.
+    Returns:
+      Tensor: The result of multiplying `tensor1` onto `tensor2`.
+    """
+    raise NotImplementedError(
+        "Backend '{}' has not implemented `broadcast_left_multiplication`."
+        .format(self.name))

tensornetwork/backends/jax/jax_backend_test.py

@@ -313,3 +313,21 @@ def test_broadcast_right_multiplication_raises():
   tensor2 = backend.randn((3, 3))
   with pytest.raises(ValueError):
     backend.broadcast_right_multiplication(tensor1, tensor2)
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
+def test_broadcast_left_multiplication(dtype):
+  backend = jax_backend.JaxBackend()
+  tensor1 = backend.randn((3,), dtype=dtype, seed=10)
+  tensor2 = backend.randn((3, 4, 2), dtype=dtype, seed=10)
+  out = backend.broadcast_left_multiplication(tensor1, tensor2)
+  np.testing.assert_allclose(out, np.reshape(tensor1, (3, 1, 1)) * tensor2)
+
+
+def test_broadcast_left_multiplication_raises():
+  dtype = np.float64
+  backend = jax_backend.JaxBackend()
+  tensor1 = backend.randn((3, 3), dtype=dtype, seed=10)
+  tensor2 = backend.randn((2, 4, 3), dtype=dtype, seed=10)
+  with pytest.raises(ValueError):
+    backend.broadcast_left_multiplication(tensor1, tensor2)

tensornetwork/backends/numpy/numpy_backend.py

@@ -405,3 +405,12 @@ def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
       raise ValueError("only order-1 tensors are allowed for `tensor2`,"
                        " found `tensor2.shape = {}`".format(tensor2.shape))
     return tensor1 * tensor2
+
+  def broadcast_left_multiplication(self, tensor1: Tensor, tensor2: Tensor):
+    if len(tensor1.shape) != 1:
+      raise ValueError("only order-1 tensors are allowed for `tensor1`,"
+                       " found `tensor1.shape = {}`".format(tensor1.shape))
+
+    t1_broadcast_shape = self.shape_concat(
+        [self.shape_tensor(tensor1), [1] * (len(tensor2.shape) - 1)], axis=-1)
+    return tensor2 * self.reshape(tensor1, t1_broadcast_shape)

tensornetwork/backends/numpy/numpy_backend_test.py

@@ -629,3 +629,21 @@ def test_broadcast_right_multiplication_raises():
   tensor2 = backend.randn((3, 3), dtype=dtype, seed=10)
   with pytest.raises(ValueError):
     backend.broadcast_right_multiplication(tensor1, tensor2)
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
+def test_broadcast_left_multiplication(dtype):
+  backend = numpy_backend.NumPyBackend()
+  tensor1 = backend.randn((3,), dtype=dtype, seed=10)
+  tensor2 = backend.randn((3, 4, 2), dtype=dtype, seed=10)
+  out = backend.broadcast_left_multiplication(tensor1, tensor2)
+  np.testing.assert_allclose(out, np.reshape(tensor1, (3, 1, 1)) * tensor2)
+
+
+def test_broadcast_left_multiplication_raises():
+  dtype = np.float64
+  backend = numpy_backend.NumPyBackend()
+  tensor1 = backend.randn((3, 3), dtype=dtype, seed=10)
+  tensor2 = backend.randn((2, 4, 3), dtype=dtype, seed=10)
+  with pytest.raises(ValueError):
+    backend.broadcast_left_multiplication(tensor1, tensor2)

tensornetwork/backends/pytorch/pytorch_backend.py

@@ -313,3 +313,12 @@ def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
           .format(tensor2.shape))
 
     return tensor1 * tensor2
+
+  def broadcast_left_multiplication(self, tensor1: Tensor, tensor2: Tensor):
+    if len(tensor1.shape) != 1:
+      raise ValueError("only order-1 tensors are allowed for `tensor1`,"
+                       " found `tensor1.shape = {}`".format(tensor1.shape))
+
+    t1_broadcast_shape = self.shape_concat(
+        [self.shape_tensor(tensor1), [1] * (len(tensor2.shape) - 1)], axis=-1)
+    return tensor2 * self.reshape(tensor1, t1_broadcast_shape)

tensornetwork/backends/pytorch/pytorch_backend_test.py

@@ -461,3 +461,21 @@ def test_broadcast_right_multiplication_raises():
   tensor2 = backend.randn((3, 3), dtype=dtype, seed=10)
   with pytest.raises(ValueError):
     backend.broadcast_right_multiplication(tensor1, tensor2)
+
+
+def test_broadcast_left_multiplication():
+  dtype = torch.float64
+  backend = pytorch_backend.PyTorchBackend()
+  tensor1 = backend.randn((3,), dtype=dtype, seed=10)
+  tensor2 = backend.randn((3, 4, 2), dtype=dtype, seed=10)
+  out = backend.broadcast_left_multiplication(tensor1, tensor2)
+  np.testing.assert_allclose(out, np.reshape(tensor1, (3, 1, 1)) * tensor2)
+
+
+def test_broadcast_left_multiplication_raises():
+  dtype = torch.float64
+  backend = pytorch_backend.PyTorchBackend()
+  tensor1 = backend.randn((3, 3), dtype=dtype, seed=10)
+  tensor2 = backend.randn((2, 4, 3), dtype=dtype, seed=10)
+  with pytest.raises(ValueError):
+    backend.broadcast_left_multiplication(tensor1, tensor2)

tensornetwork/backends/shell/shell_backend.py

@@ -332,3 +332,15 @@ def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
       shape2 = tuple([1] * (len(tensor1.shape) - len(shape2))) + shape2
     shape = tuple([max([s1, s2]) for s1, s2 in zip(tensor1.shape, shape2)])
     return ShellTensor(shape)
+
+  def broadcast_left_multiplication(self, tensor1: Tensor, tensor2: Tensor):
+    if len(tensor1.shape) != 1:
+      raise ValueError(
+          "only order-1 tensors are allowed for `tensor1`, found `tensor1.shape = {}`"
+          .format(tensor1.shape))
+
+    shape1 = tuple(tensor1.shape)
+    if len(shape1) < len(tensor2.shape):
+      shape1 = shape1 + tuple([1] * (len(tensor2.shape) - len(shape1)))
+    shape = tuple([max([s1, s2]) for s1, s2 in zip(tensor2.shape, shape1)])
+    return ShellTensor(shape)

tensornetwork/backends/shell/shell_backend_test.py

@@ -293,3 +293,19 @@ def test_broadcast_right_multiplication_raises():
   tensor2 = backend.randn((3, 3))
   with pytest.raises(ValueError):
     backend.broadcast_right_multiplication(tensor1, tensor2)
+
+
+def test_broadcast_left_multiplication():
+  backend = shell_backend.ShellBackend()
+  tensor1 = backend.randn((3,))
+  tensor2 = backend.randn((3, 4, 2))
+  out = backend.broadcast_left_multiplication(tensor1, tensor2)
+  np.testing.assert_allclose(out.shape, [3, 4, 2])
+
+
+def test_broadcast_left_multiplication_raises():
+  backend = shell_backend.ShellBackend()
+  tensor1 = backend.randn((3, 3))
+  tensor2 = backend.randn((3, 4, 2))
+  with pytest.raises(ValueError):
+    backend.broadcast_left_multiplication(tensor1, tensor2)

tensornetwork/backends/symmetric/symmetric_backend.py

@@ -196,9 +196,11 @@ def broadcast_right_multiplication(self, tensor1: Tensor, tensor2: Tensor):
     if tensor2.ndim != 1:
       raise ValueError("only order-1 tensors are allowed for `tensor2`,"
                        " found `tensor2.shape = {}`".format(tensor2.shape))
-    shape1 = tensor1.shape
-    tmp = self.reshape(tensor1,
-                       (numpy.prod(shape1[:tensor1.ndim - 1]), shape1[-1]))
-    #NOTE (mganahl): we use mulitplication by diagonal matrix here
-    return self.reshape(
-        self.tensordot(tmp, self.diag(tensor2), ([1], [0])), shape1)
+    return self.tensordot(tensor1, self.diag(tensor2),
+                          ([len(tensor1.shape) - 1], [0]))
+
+  def broadcast_left_multiplication(self, tensor1: Tensor, tensor2: Tensor):
+    if len(tensor1.shape) != 1:
+      raise ValueError("only order-1 tensors are allowed for `tensor1`,"
+                       " found `tensor1.shape = {}`".format(tensor1.shape))
+    return self.tensordot(self.diag(tensor1), tensor2, ([1], [0]))

tensornetwork/backends/symmetric/symmetric_backend_test.py

@@ -739,3 +739,45 @@ def test_broadcast_right_multiplication_raises():
   tensor2 = ChargeArray.random(indices=indices)
   with pytest.raises(ValueError):
     backend.broadcast_right_multiplication(tensor1, tensor2)
+
+
+@pytest.mark.parametrize("dtype", [np.float64, np.complex128])
+@pytest.mark.parametrize("num_charges", [1, 2])
+def test_broadcast_left_multiplication(dtype, num_charges):
+  np.random.seed(10)
+  backend = symmetric_backend.SymmetricBackend()
+  Ds = [10, 30, 24]
+  R = len(Ds)
+  indices = [
+      Index(
+          BaseCharge(
+              np.random.randint(-5, 6, (num_charges, Ds[n])),
+              charge_types=[U1Charge] * num_charges), False) for n in range(R)
+  ]
+
+  tensor1 = ChargeArray.random(indices=[indices[0]], dtype=dtype)
+  tensor2 = backend.randn(indices, dtype=dtype)
+  t1dense = tensor1.todense()
+  t2dense = tensor2.todense()
+  out = backend.broadcast_left_multiplication(tensor1, tensor2)
+  dense = np.reshape(t1dense, (10, 1, 1)) * t2dense
+  np.testing.assert_allclose(out.todense(), dense)
+
+
+def test_broadcast_left_multiplication_raises():
+  np.random.seed(10)
+  backend = symmetric_backend.SymmetricBackend()
+  num_charges = 1
+  Ds = [10, 30, 24]
+  R = len(Ds)
+  indices = [
+      Index(
+          BaseCharge(
+              np.random.randint(-5, 6, (num_charges, Ds[n])),
+              charge_types=[U1Charge] * num_charges), False) for n in range(R)
+  ]
+
+  tensor1 = ChargeArray.random(indices=indices)
+  tensor2 = backend.randn(indices)
+  with pytest.raises(ValueError):
+    backend.broadcast_left_multiplication(tensor1, tensor2)