Repository: cornellius-gp/gpytorch
Issue: [Bug] ConstantMulLazyTensor recomputes its own root

# 🐛 Bug

`ConstantMulLazyTensor` recomputes its own root decomposition even when the component lazy tensor already has its own root. This is quite slow because it means more root decompositions are called than are strictly necessary.

## To reproduce

** Code snippet to reproduce **
```python
import torch
from gpytorch.lazy import lazify, RootLazyTensor

a = torch.randn(50, 10)
lazy_square_a = RootLazyTensor(lazify(a))
(4. * lazy_square_a).root_decomposition().root.evaluate() # this is not 2 * a.
```

## Expected Behavior

At least for positive constants, this should return 2 * a.

## System information

**Please complete the following information:**
gpytorch master
pytorch 1.7.0
mac os

## Additional context
I'll make a PR in a minute.


Changed files:
gpytorch/distributions/multivariate_normal.py, gpytorch/lazy/root_lazy_tensor.py, gpytorch/lazy/constant_mul_lazy_tensor.py

diff --git a/gpytorch/distributions/multivariate_normal.py b/gpytorch/distributions/multivariate_normal.py
index 48d360e0..de710ed9 100644
--- a/gpytorch/distributions/multivariate_normal.py
+++ b/gpytorch/distributions/multivariate_normal.py
@@ -9,7 +9,7 @@ from torch.distributions.kl import register_kl
 from torch.distributions.utils import _standard_normal, lazy_property
 
 from .. import settings
-from ..lazy import DiagLazyTensor, LazyTensor, delazify, lazify
+from ..lazy import DiagLazyTensor, LazyTensor, RootLazyTensor, delazify, lazify
 from ..utils.broadcasting import _mul_broadcast_shape
 from ..utils.warnings import NumericalWarning
 from .distribution import Distribution
@@ -42,7 +42,17 @@ class MultivariateNormal(TMultivariateNormal, Distribution):
             self.__unbroadcasted_scale_tril = None
             self._validate_args = validate_args
             batch_shape = _mul_broadcast_shape(self.loc.shape[:-1], covariance_matrix.shape[:-2])
-            event_shape = self.loc.shape[-1:]
+            if not isinstance(self._covar, RootLazyTensor):
+                should_have_separate_event = False
+            else:
+                should_have_separate_event = True
+            if should_have_separate_event:
+                should_have_separate_event = self._covar.root.shape[-1] > self._covar.root.shape[-2]
+            if not should_have_separate_event:
+                event_shape = self.loc.shape[-1:]
+            else:
+                event_shape = self._covar.root.shape[-1:]
+
             # TODO: Integrate argument validation for LazyTensors into torch.distribution validation logic
             super(TMultivariateNormal, self).__init__(batch_shape, event_shape, validate_args=False)
         else:
@@ -154,8 +164,13 @@ class MultivariateNormal(TMultivariateNormal, Distribution):
             res = res.view(sample_shape + self.loc.shape)
 
         else:
+            covar_root = covar.root_decomposition().root
+
             # Make sure that the base samples agree with the distribution
-            if self.loc.shape != base_samples.shape[-self.loc.dim() :]:
+            if (
+                self.loc.shape != base_samples.shape[-self.loc.dim() :]
+                and covar_root.shape[-1] < base_samples.shape[-1]
+            ):
                 raise RuntimeError(
                     "The size of base_samples (minus sample shape dimensions) should agree with the size "
                     "of self.loc. Expected ...{} but got {}".format(self.loc.shape, base_samples.shape)
@@ -166,16 +181,16 @@ class MultivariateNormal(TMultivariateNormal, Distribution):
 
             # Reshape samples to be batch_size x num_dim x num_samples
             # or num_bim x num_samples
-            base_samples = base_samples.view(-1, *self.loc.shape)
+            base_samples = base_samples.view(-1, *self.loc.shape[:-1], covar_root.shape[-1])
             base_samples = base_samples.permute(*range(1, self.loc.dim() + 1), 0)
 
             # Now reparameterize those base samples
-            covar_root = covar.root_decomposition().root
             # If necessary, adjust base_samples for rank of root decomposition
             if covar_root.shape[-1] < base_samples.shape[-2]:
                 base_samples = base_samples[..., : covar_root.shape[-1], :]
             elif covar_root.shape[-1] > base_samples.shape[-2]:
-                raise RuntimeError("Incompatible dimension of `base_samples`")
+                # raise RuntimeError("Incompatible dimension of `base_samples`")
+                covar_root = covar_root.transpose(-2, -1)
             res = covar_root.matmul(base_samples) + self.loc.unsqueeze(-1)
 
             # Permute and reshape new samples to be original size
diff --git a/gpytorch/lazy/constant_mul_lazy_tensor.py b/gpytorch/lazy/constant_mul_lazy_tensor.py
index 9ca85079..bd98f995 100644
--- a/gpytorch/lazy/constant_mul_lazy_tensor.py
+++ b/gpytorch/lazy/constant_mul_lazy_tensor.py
@@ -1,9 +1,11 @@
 #!/usr/bin/env python3
+from typing import Optional
 
 import torch
 
 from ..utils.memoize import cached
 from .lazy_tensor import LazyTensor
+from .root_lazy_tensor import RootLazyTensor
 
 
 class ConstantMulLazyTensor(LazyTensor):
@@ -153,3 +155,11 @@ class ConstantMulLazyTensor(LazyTensor):
     def evaluate(self):
         res = self.base_lazy_tensor.evaluate()
         return res * self.expanded_constant
+
+    @cached(name="root_decomposition")
+    def root_decomposition(self, method: Optional[str] = None):
+        if torch.all(self._constant >= 0):
+            base_root = self.base_lazy_tensor.root_decomposition(method=method).root
+            return RootLazyTensor(ConstantMulLazyTensor(base_root, self._constant ** 0.5))
+
+        return super().root_decomposition(method=method)
diff --git a/gpytorch/lazy/root_lazy_tensor.py b/gpytorch/lazy/root_lazy_tensor.py
index 0373c5ad..1ae423f7 100644
--- a/gpytorch/lazy/root_lazy_tensor.py
+++ b/gpytorch/lazy/root_lazy_tensor.py
@@ -66,7 +66,7 @@ class RootLazyTensor(LazyTensor):
         # Matrix is symmetric
         return self._matmul(rhs)
 
-    def root_decomposition(self):
+    def root_decomposition(self, method=None):
         return self
 
     def _root_decomposition(self):