Repository: cornellius-gp/gpytorch

Issue:
[Feature Request] Backward hooks in Exact GPs
# 🚀 Feature Request

In some cases we want to use larger models that are composed of different models, e.g. for modeling noise level with a separate GP in a heteroskedastic likelihood. In this case the noise model needs to be evaluated in eval mode while the other one is in train mode. For this to work properly, we need to make sure the test caches get cleared after backpropagating through them. 

A specific implementation for the noise models was done in https://github.com/cornellius-gp/gpytorch/pull/870, but this would be more generally useful and we should consider making this part of `ExactGP`. 



Discussion:

Resolving patch:
diff --git a/gpytorch/likelihoods/noise_models.py b/gpytorch/likelihoods/noise_models.py
--- a/gpytorch/likelihoods/noise_models.py
+++ b/gpytorch/likelihoods/noise_models.py
@@ -130,8 +130,6 @@ def __init__(self, noise_model, noise_indices=None, noise_constraint=None):
         self.noise_model = noise_model
         self._noise_constraint = noise_constraint
         self._noise_indices = noise_indices
-        # make sure to clear test caches after backpropagating through them
-        self.register_backward_hook(clear_cache_hook)
 
     def forward(
         self,
@@ -190,10 +188,3 @@ def forward(
     def _apply(self, fn):
         self.noise = fn(self.noise)
         return super(FixedGaussianNoise, self)._apply(fn)
-
-
-def clear_cache_hook(module, grad_input, grad_output) -> None:
-    try:
-        module.noise_model.prediction_strategy = None
-    except AttributeError:
-        pass
diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
--- a/gpytorch/models/exact_prediction_strategies.py
+++ b/gpytorch/models/exact_prediction_strategies.py
@@ -19,6 +19,11 @@
 )
 from ..utils.interpolation import left_interp, left_t_interp
 from ..utils.memoize import cached, add_to_cache
+import functools
+
+
+def clear_cache_hook(module, *args, **kwargs):
+    module._memoize_cache = {}
 
 
 def prediction_strategy(train_inputs, train_prior_dist, train_labels, likelihood):
@@ -69,10 +74,16 @@ def _exact_predictive_covar_inv_quad_form_cache(self, train_train_covar_inv_root
         Returns
             - A precomputed cache
         """
+        res = train_train_covar_inv_root
         if settings.detach_test_caches.on():
-            return train_train_covar_inv_root.detach()
-        else:
-            return train_train_covar_inv_root
+            res = res.detach()
+
+        if res.grad_fn is not None:
+            wrapper = functools.partial(clear_cache_hook, self)
+            functools.update_wrapper(wrapper, clear_cache_hook)
+            res.grad_fn.register_hook(wrapper)
+
+        return res
 
     def _exact_predictive_covar_inv_quad_form_root(self, precomputed_cache, test_train_covar):
         """
@@ -257,9 +268,14 @@ def mean_cache(self):
         mean_cache = train_train_covar.inv_matmul(train_labels_offset).squeeze(-1)
 
         if settings.detach_test_caches.on():
-            return mean_cache.detach()
-        else:
-            return mean_cache
+            mean_cache = mean_cache.detach()
+
+        if mean_cache.grad_fn is not None:
+            wrapper = functools.partial(clear_cache_hook, self)
+            functools.update_wrapper(wrapper, clear_cache_hook)
+            mean_cache.grad_fn.register_hook(wrapper)
+
+        return mean_cache
 
     @property
     def num_train(self):
@@ -302,6 +318,7 @@ def exact_predictive_mean(self, test_mean, test_train_covar):
         # GP, and using addmv requires you to delazify test_train_covar, which is obviously a huge no-no!
         res = (test_train_covar @ self.mean_cache.unsqueeze(-1)).squeeze(-1)
         res = res + test_mean
+
         return res
 
     def exact_predictive_covar(self, test_test_covar, test_train_covar):


Test patch: