Repository: cornellius-gp/gpytorch
Pull request #886 (2019-09-28 22:28:05)

Add ability to add torch.Tensor into SumLazyTensor
Provides a partial fix for #879. This PR works provided that the torch.tensor is the tensor that needs to be broadcasted and NOT the lazyTensor itself. 

Resolving lazyTensor broadcasting itself is a bigger issue, given the following code: 
```
res = lazify(torch.randn(30, 1)) + torch.randn(30, 30)
print(res.shape) # 30, 1 not 30, 30
```

Will attempt to look into more detail when I get another chance.



Patch:
diff --git a/.isort.cfg b/.isort.cfg
index 92b18f67..3e8d8aaf 100644
--- a/.isort.cfg
+++ b/.isort.cfg
@@ -1,2 +1,2 @@
 [settings]
-known_third_party = matplotlib,numpy,pyro,setuptools,sphinx_rtd_theme,torch
+known_third_party = gpytorch,matplotlib,numpy,pyro,setuptools,sphinx_rtd_theme,torch
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 1e37ac58..89c433f4 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -1611,11 +1611,17 @@ class LazyTensor(ABC):
         from .zero_lazy_tensor import ZeroLazyTensor
         from .diag_lazy_tensor import DiagLazyTensor
         from .added_diag_lazy_tensor import AddedDiagLazyTensor
+        from .non_lazy_tensor import lazify
+        from torch import Tensor
 
         if isinstance(other, ZeroLazyTensor):
             return self
         elif isinstance(other, DiagLazyTensor):
             return AddedDiagLazyTensor(self, other)
+        elif isinstance(other, Tensor):
+            other = lazify(other)
+            shape = _mul_broadcast_shape(self.shape, other.shape)
+            return SumLazyTensor(self.expand(shape), other.expand(shape))
         else:
             return SumLazyTensor(self, other)
 
diff --git a/gpytorch/lazy/non_lazy_tensor.py b/gpytorch/lazy/non_lazy_tensor.py
index c4a0dfed..225ca490 100644
--- a/gpytorch/lazy/non_lazy_tensor.py
+++ b/gpytorch/lazy/non_lazy_tensor.py
@@ -72,6 +72,8 @@ class NonLazyTensor(LazyTensor):
     def __add__(self, other):
         if isinstance(other, NonLazyTensor):
             return NonLazyTensor(self.tensor + other.tensor)
+        elif isinstance(other, torch.Tensor):
+            return NonLazyTensor(self.tensor + other)
         else:
             return super(NonLazyTensor, self).__add__(other)
 
diff --git a/gpytorch/lazy/sum_lazy_tensor.py b/gpytorch/lazy/sum_lazy_tensor.py
index ec01e5d3..f083b21b 100644
--- a/gpytorch/lazy/sum_lazy_tensor.py
+++ b/gpytorch/lazy/sum_lazy_tensor.py
@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+from torch import Tensor
 
 from ..utils.broadcasting import _mul_broadcast_shape
 from ..utils.memoize import cached
@@ -6,15 +7,17 @@ from .lazy_tensor import LazyTensor
 from .non_lazy_tensor import lazify
 from .zero_lazy_tensor import ZeroLazyTensor
 
+# from .broadcasted_lazy_tensor import BroadcastedLazyTensor
+
 
 class SumLazyTensor(LazyTensor):
     def __init__(self, *lazy_tensors, **kwargs):
-        lazy_tensors = list(lazy_tensors)
-        for i, lazy_tensor in enumerate(lazy_tensors):
-            try:
-                lazy_tensors[i] = lazify(lazy_tensor)
-            except TypeError:
-                raise TypeError("All arguments of a SumLazyTensor should be LazyTensors or Tensors")
+        try:
+            lazy_tensors = tuple(lazify(lt) for lt in lazy_tensors)
+        except TypeError:
+            raise TypeError("All arguments of a SumLazyTensor should be LazyTensors or Tensors")
+        # batch_shape = _mul_broadcast_shape(*[lt.shape for lt in lazy_tensors])
+        # lazy_tensors = tuple(lt.expand(batch_shape) for lt in lazy_tensors)
         super(SumLazyTensor, self).__init__(*lazy_tensors, **kwargs)
 
         self.lazy_tensors = lazy_tensors
@@ -68,6 +71,22 @@ class SumLazyTensor(LazyTensor):
             return SumLazyTensor(*(list(self.lazy_tensors) + list(other.lazy_tensors)))
         elif isinstance(other, LazyTensor):
             return SumLazyTensor(*(list(self.lazy_tensors) + [other]))
+        elif isinstance(other, Tensor):
+            # get broadcast shape, assuming mul broadcasting the same as add broadcasting
+            broadcasted_shape = _mul_broadcast_shape(self.shape, other.shape)
+
+            # lazify + broadcast other
+            broadcasted_other = lazify(other.expand(broadcasted_shape))
+
+            # update the lazy tensors' shape as well
+            if broadcasted_shape != self.shape:
+                broadcasted_lts = [
+                    lt.expand(*broadcasted_shape, 1).squeeze(-1).transpose(-1, -2) for lt in self.lazy_tensors
+                ]
+            else:
+                broadcasted_lts = list(self.lazy_tensors)
+
+            return SumLazyTensor(*(broadcasted_lts + [broadcasted_other]))
         else:
             raise AttributeError("other must be a LazyTensor")
 


Test patch:
diff --git a/test/lazy/test_sum_lazy_tensor.py b/test/lazy/test_sum_lazy_tensor.py
index d9e8dfda..0b5842bb 100644
--- a/test/lazy/test_sum_lazy_tensor.py
+++ b/test/lazy/test_sum_lazy_tensor.py
@@ -4,7 +4,7 @@ import unittest
 
 import torch
 
-from gpytorch.lazy import ToeplitzLazyTensor
+from gpytorch.lazy import ToeplitzLazyTensor, lazify
 from gpytorch.test.lazy_tensor_test_case import LazyTensorTestCase
 
 
@@ -45,11 +45,11 @@ class TestSumLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
 
     def create_lazy_tensor(self):
         c1 = torch.tensor(
-            [[[2, 0.5, 0, 0], [5, 1, 2, 0]], [[2, 0.5, 0, 0], [5, 1, 2, 0]]], dtype=torch.float, requires_grad=True
+            [[[2, 0.5, 0, 0], [5, 1, 2, 0]], [[2, 0.5, 0, 0], [5, 1, 2, 0]]], dtype=torch.float, requires_grad=True,
         )
         t1 = ToeplitzLazyTensor(c1)
         c2 = torch.tensor(
-            [[[2, 0.5, 0, 0], [5, 1, 2, 0]], [[2, 0.5, 0, 0], [6, 0, 1, -1]]], dtype=torch.float, requires_grad=True
+            [[[2, 0.5, 0, 0], [5, 1, 2, 0]], [[2, 0.5, 0, 0], [6, 0, 1, -1]]], dtype=torch.float, requires_grad=True,
         )
         t2 = ToeplitzLazyTensor(c2)
         return t1 + t2
@@ -59,5 +59,43 @@ class TestSumLazyTensorMultiBatch(LazyTensorTestCase, unittest.TestCase):
         return sum(tensors)
 
 
+class TestSumLazyTensorBroadcasting(unittest.TestCase):
+    def test_broadcast_same_shape(self):
+        test1 = lazify(torch.randn(30, 30))
+
+        test2 = torch.randn(30, 30)
+        res = test1 + test2
+        final_res = res + test2
+
+        torch_res = res.evaluate() + test2
+
+        self.assertEqual(final_res.shape, torch_res.shape)
+        self.assertEqual((final_res.evaluate() - torch_res).sum(), 0.0)
+
+    def test_broadcast_tensor_shape(self):
+        test1 = lazify(torch.randn(30, 30))
+
+        test2 = torch.randn(30, 1)
+        res = test1 + test2
+        final_res = res + test2
+
+        torch_res = res.evaluate() + test2
+
+        self.assertEqual(final_res.shape, torch_res.shape)
+        self.assertEqual((final_res.evaluate() - torch_res).sum(), 0.0)
+
+    def test_broadcast_lazy_shape(self):
+        test1 = lazify(torch.randn(30, 1))
+
+        test2 = torch.randn(30, 30)
+        res = test1 + test2
+        final_res = res + test2
+
+        torch_res = res.evaluate() + test2
+
+        self.assertEqual(final_res.shape, torch_res.shape)
+        self.assertEqual((final_res.evaluate() - torch_res).sum(), 0.0)
+
+
 if __name__ == "__main__":
     unittest.main()