Repository: cornellius-gp/gpytorch

Issue:
[Bug] LazyEvaluatedKernelTensor does not require gradient if inputs do not require gradient
# 🐛 Bug

Not entirely sure if this is intended behavior, but it does produce some downstream problems. 
If the input to `kernel(x)` doesn't require a gradient `LazyEvaluatedKernelTensor` does not portray as requiring a gradient, even though it does. Crucially `kernel(x).evaluate_kernel()` does. 

As a result, it causes some weirdness downstream in objects that can manipulate lazy kernel tensors themselves:

## To reproduce

** Code snippet to reproduce **
```python
import torch
from gpytorch.kernels import MaternKernel
from gpytorch.lazy import KroneckerProductLazyTensor

torch.random.manual_seed(3)

kernels = [MaternKernel()] * 3
x = [torch.randn(5)] * 3

kernels[0](x[0]).requires_grad # False

# breaking behavior
kplt = KroneckerProductLazyTensor(*[k(a) for k, a in zip(kernels, x)])

rhs = torch.randn(kplt.shape[-1])
print(kplt.matmul(rhs).requires_grad) # True, as expected

kdiag = kplt.diag()
print(kdiag.requires_grad) # False ??

# the fix is that this works fine
kplt = KroneckerProductLazyTensor(*[k(a).evaluate_kernel() for k, a in zip(kernels, x)])
kdiag = kplt.diag()
print(kdiag.requires_grad) # True
```


## Expected Behavior

LazyKernelEvaluatedTensor should probably require a gradient to prevent this sort of behavior from happening.

## System information

**Please complete the following information:**
gpytorch master

## Additional context

I'm trying to implement structured variational MTGP models with non-gaussian likelihoods and am struggling to get reasonable gradients, noticed while debugging.



Discussion:
```python
# breaking behavior
kplt = KroneckerProductLazyTensor(*[k(a) for k, a in zip(kernels, x)])

kdiag = kplt.diag()
print(kdiag.requires_grad) # False ??
```

This is actually correct behavior for this specific script: when calling `.diag` of a stationary kernel, we populate the distance function with a zeros vector (which doesn't require grad). See https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/kernels/kernel.py#L312

However, in the more general case (e.g. replace `MaternKernel` with `LinearKernel` in your example), then yes there is breaking behavior. I'm putting up a PR to fix it.


Resolving patch:
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
--- a/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
+++ b/gpytorch/lazy/lazy_evaluated_kernel_tensor.py
@@ -37,6 +37,16 @@ def dtype(self):
     def device(self):
         return self.x1.device
 
+    @property
+    def requires_grad(self):
+        return super().requires_grad or any(param.requires_grad for param in self.kernel.parameters())
+
+    def _set_requires_grad(self, val):
+        super()._set_requires_grad(val)
+        # The behavior that differs from the base LazyTensor setter
+        for param in self.kernel.parameters():
+            param.requires_grad_(val)
+
     def _expand_batch(self, batch_shape):
         return self.evaluate_kernel()._expand_batch(batch_shape)
 
@@ -91,8 +101,8 @@ def _getitem(self, row_index, col_index, *batch_indices):
         except IndexError:
             if any(not isinstance(bi, slice) for bi in batch_indices):
                 raise RuntimeError(
-                    f"Attempting to tensor index a non-batch matrix's batch dimensions. "
-                    "Got batch index {batch_indices} but my shape was {self.shape}"
+                    "Attempting to tensor index a non-batch matrix's batch dimensions. "
+                    f"Got batch index {batch_indices} but my shape was {self.shape}"
                 )
             x1 = x1.expand(*([1] * (len(batch_indices) - self.x1.dim() + 2)), *self.x1.shape)
             x1 = x1[(*batch_indices, row_index, dim_index)]
@@ -105,8 +115,8 @@ def _getitem(self, row_index, col_index, *batch_indices):
         except IndexError:
             if any([not isinstance(bi, slice) for bi in batch_indices]):
                 raise RuntimeError(
-                    f"Attempting to tensor index a non-batch matrix's batch dimensions. "
-                    "Got batch index {batch_indices} but my shape was {self.shape}"
+                    "Attempting to tensor index a non-batch matrix's batch dimensions. "
+                    f"Got batch index {batch_indices} but my shape was {self.shape}"
                 )
             x2 = x2.expand(*([1] * (len(batch_indices) - self.x1.dim() + 2)), *self.x2.shape)
             x2 = x2[(*batch_indices, col_index, dim_index)]
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -1529,22 +1529,28 @@ def requires_grad(self):
             if hasattr(arg, "requires_grad")
         )
 
-    @requires_grad.setter
-    def requires_grad(self, val):
+    def _set_requires_grad(self, val):
+        # Note: subclasses should overwrite this method, not the requires_grad.setter
         for arg in self._args:
             if hasattr(arg, "requires_grad"):
                 if arg.dtype in (torch.float, torch.double, torch.half):
-                    arg.requires_grad = val
+                    arg.requires_grad_(val)
         for arg in self._kwargs.values():
             if hasattr(arg, "requires_grad"):
-                arg.requires_grad = val
+                arg.requires_grad_(val)
+
+    @requires_grad.setter
+    def requires_grad(self, val):
+        # Note: subclasses cannot overwrite this method
+        # To change the setter behavior, overwrite the _set_requires_grad method instead
+        self._set_requires_grad(val)
 
     def requires_grad_(self, val):
         """
         Sets `requires_grad=val` on all the Tensors that make up the LazyTensor
         This is an inplace operation.
         """
-        self.requires_grad = val
+        self._set_requires_grad(val)
         return self
 
     @cached(name="diagonalization")


Test patch: