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.



Discussion:

Resolving patch:
diff --git a/gpytorch/distributions/multivariate_normal.py b/gpytorch/distributions/multivariate_normal.py
--- a/gpytorch/distributions/multivariate_normal.py
+++ b/gpytorch/distributions/multivariate_normal.py
@@ -9,7 +9,7 @@
 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 @@ def __init__(self, mean, covariance_matrix, validate_args=False):
             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 @@ def rsample(self, sample_shape=torch.Size(), base_samples=None):
             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 @@ def rsample(self, sample_shape=torch.Size(), base_samples=None):
 
             # 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
--- 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 @@ def diag(self):
     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
--- a/gpytorch/lazy/root_lazy_tensor.py
+++ b/gpytorch/lazy/root_lazy_tensor.py
@@ -66,7 +66,7 @@ def _t_matmul(self, rhs):
         # Matrix is symmetric
         return self._matmul(rhs)
 
-    def root_decomposition(self):
+    def root_decomposition(self, method=None):
         return self
 
     def _root_decomposition(self):


Test patch: