Repository: cornellius-gp/gpytorch
Issue: [Bug] ConstantMulLazyTensor can only add batched tensors

# 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

ConstantMulLazyTensor attempts to expand `_constant` according to batch dimensions which fails when there is no batch dimension.

## To reproduce

** Code snippet to reproduce **
```python
import gpytorch
import torch

tlt = gpytorch.lazy.ToeplitzLazyTensor(torch.randn(100).abs())
cmtlt = torch.ones(1) * tlt
cmtlt + torch.randn(100, 100)

#more specifically
cmtlt._constant.expand(*torch.Size([])) # fails
```

** Stack trace/error message **
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-16-0d626e17adaf> in <module>
----> 1 cmtlt + torch.randn(100, 100)

~/Documents/GitHub/gpytorch/gpytorch/lazy/lazy_tensor.py in __add__(self, other)
   1625             other = lazify(other)
   1626             shape = _mul_broadcast_shape(self.shape, other.shape)
-> 1627             return SumLazyTensor(self.expand(shape), other.expand(shape))
   1628         else:
   1629             return SumLazyTensor(self, other)

~/Documents/GitHub/gpytorch/gpytorch/lazy/lazy_tensor.py in expand(self, *sizes)
    857             raise RuntimeError("Invalid arguments {} to expand.".format(sizes))
    858 
--> 859         res = self._expand_batch(batch_shape=shape[:-2])
    860         return res
    861 

~/Documents/GitHub/gpytorch/gpytorch/lazy/constant_mul_lazy_tensor.py in _expand_batch(self, batch_shape)
     70     def _expand_batch(self, batch_shape):
---> 71         return self.__class__(self.base_lazy_tensor._expand_batch(batch_shape), self._constant.expand(*batch_shape))
     72 
     73     def _get_indices(self, row_index, col_index, *batch_indices):

TypeError: expand() missing 1 required positional arguments: "size"
```

## Expected Behavior

It should return a SumLazyTensor of shape 100 x 100.

## System information

**Please complete the following information:**
gpytorch 1.0.1 (code hasn't changed in 1.1.0)
pytorch 1.4.0
MacOS

## Additional context
The fix is to force the tensor that is being added to have a batch dimension.
```python
cmtlt + torch.randn(1, 100, 100)
```



Changed files:
gpytorch/lazy/lazy_tensor.py, gpytorch/lazy/constant_mul_lazy_tensor.py, gpytorch/lazy/sum_lazy_tensor.py, gpytorch/lazy/cat_lazy_tensor.py, gpytorch/test/lazy_tensor_test_case.py

diff --git a/gpytorch/lazy/cat_lazy_tensor.py b/gpytorch/lazy/cat_lazy_tensor.py
index 989a1573..43b430da 100644
--- a/gpytorch/lazy/cat_lazy_tensor.py
+++ b/gpytorch/lazy/cat_lazy_tensor.py
@@ -126,7 +126,20 @@ class CatLazyTensor(LazyTensor):
             )
 
     def _expand_batch(self, batch_shape):
-        lazy_tensors = [lazy_tensor._expand_batch(batch_shape) for lazy_tensor in self.lazy_tensors]
+        batch_dim = self.cat_dim + 2
+        if batch_dim < 0:
+            if batch_shape[batch_dim] != self.batch_shape[batch_dim]:
+                raise RuntimeError(
+                    f"Trying to expand a CatLazyTensor in dimension {self.cat_dim}, but this is the concatenated "
+                    f"dimension.\\nCurrent shape: {self.shape} - expanded shape: {batch_shape + self.matrix_shape}."
+                )
+            lazy_tensors = []
+            for lazy_tensor in self.lazy_tensors:
+                sub_batch_shape = list(batch_shape).copy()
+                sub_batch_shape[batch_dim] = lazy_tensor.shape[self.cat_dim]
+                lazy_tensors.append(lazy_tensor._expand_batch(sub_batch_shape))
+        else:
+            lazy_tensors = [lazy_tensor._expand_batch(batch_shape) for lazy_tensor in self.lazy_tensors]
         res = self.__class__(*lazy_tensors, dim=self.cat_dim, output_device=self.output_device)
         return res
 
diff --git a/gpytorch/lazy/constant_mul_lazy_tensor.py b/gpytorch/lazy/constant_mul_lazy_tensor.py
index afa56387..9ca85079 100644
--- a/gpytorch/lazy/constant_mul_lazy_tensor.py
+++ b/gpytorch/lazy/constant_mul_lazy_tensor.py
@@ -68,7 +68,10 @@ class ConstantMulLazyTensor(LazyTensor):
         return res * self._constant.unsqueeze(-1)
 
     def _expand_batch(self, batch_shape):
-        return self.__class__(self.base_lazy_tensor._expand_batch(batch_shape), self._constant.expand(*batch_shape))
+        return self.__class__(
+            self.base_lazy_tensor._expand_batch(batch_shape),
+            self._constant.expand(*batch_shape) if len(batch_shape) else self._constant,
+        )
 
     def _get_indices(self, row_index, col_index, *batch_indices):
         # NOTE TO FUTURE SELF:
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 2278202d..485d41c3 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -1624,7 +1624,9 @@ class LazyTensor(ABC):
         elif isinstance(other, Tensor):
             other = lazify(other)
             shape = _mul_broadcast_shape(self.shape, other.shape)
-            return SumLazyTensor(self.expand(shape), other.expand(shape))
+            new_self = self if self.shape == shape else self._expand_batch(shape[:-2])
+            new_other = other if other.shape == shape else other._expand_batch(shape[:-2])
+            return SumLazyTensor(new_self, new_other)
         else:
             return SumLazyTensor(self, other)
 
diff --git a/gpytorch/lazy/sum_lazy_tensor.py b/gpytorch/lazy/sum_lazy_tensor.py
index 3c6c4c7e..79e263f4 100644
--- a/gpytorch/lazy/sum_lazy_tensor.py
+++ b/gpytorch/lazy/sum_lazy_tensor.py
@@ -81,14 +81,9 @@ class SumLazyTensor(LazyTensor):
             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]))
+            new_self = self if broadcasted_shape == self.shape else self._expand_batch(broadcasted_shape[:-2])
+
+            return SumLazyTensor(*(list(new_self.lazy_tensors) + [broadcasted_other]))
         else:
             raise AttributeError("other must be a LazyTensor")
 
diff --git a/gpytorch/test/lazy_tensor_test_case.py b/gpytorch/test/lazy_tensor_test_case.py
index d1a491cb..3fa08aa6 100644
--- a/gpytorch/test/lazy_tensor_test_case.py
+++ b/gpytorch/test/lazy_tensor_test_case.py
@@ -45,6 +45,19 @@ class RectangularLazyTensorTestCase(BaseTestCase):
             if arg_copy.grad is not None:
                 self.assertAllClose(arg.grad, arg_copy.grad, rtol=1e-3)
 
+    def test_add(self):
+        lazy_tensor = self.create_lazy_tensor()
+        evaluated = self.evaluate_lazy_tensor(lazy_tensor)
+
+        rhs = torch.randn(lazy_tensor.shape)
+        self.assertAllClose((lazy_tensor + rhs).evaluate(), evaluated + rhs)
+
+        rhs = torch.randn(lazy_tensor.matrix_shape)
+        self.assertAllClose((lazy_tensor + rhs).evaluate(), evaluated + rhs)
+
+        rhs = torch.randn(2, *lazy_tensor.shape)
+        self.assertAllClose((lazy_tensor + rhs).evaluate(), evaluated + rhs)
+
     def test_matmul_vec(self):
         lazy_tensor = self.create_lazy_tensor()
         rhs = torch.randn(lazy_tensor.size(-1))