We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Code in variational_strategy.py
# 🐛 Bug

Dear Gpytorch Developers

Here is a possible issue with variational_strategy.py

I am wondering whether it is a special consideration or something missing.  According to the formula in Line 121, it seems an L is missing in front of self.prior_distribution.mean in Line 115.   I understand variable inducing_values has already included L.inv() according to whitening of the inducing posterior, but this is not the case for the prior.   It is fortunate that self.prior_distribution.mean is always 0, so no errors.

Thank you for the great package

J.

## To reproduce

** Code snippet to reproduce **
```python
# Your code goes here
# Please make sure it does not require any external dependencies (other than PyTorch!)
# (We much prefer small snippets rather than links to existing libraries!)
```

** Stack trace/error message **
```
// Paste the bad output here!
```

## Expected Behavior

<!-- A clear and concise description of what you expected to happen. -->

## System information

**Please complete the following information:**
- <!-- GPyTorch Version (run `print(gpytorch.__version__)` -->
- <!-- PyTorch Version (run `print(torch.__version__)` -->
- <!-- Computer OS -->

## Additional context
Add any other context about the problem here.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `gpytorch/variational/variational_strategy.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 import warnings
4 
5 import torch
6 
7 from .. import settings
8 from ..distributions import MultivariateNormal
9 from ..lazy import DiagLazyTensor, MatmulLazyTensor, RootLazyTensor, SumLazyTensor, TriangularLazyTensor, delazify
10 from ..settings import trace_mode
11 from ..utils.cholesky import psd_safe_cholesky
12 from ..utils.errors import CachingError
13 from ..utils.memoize import cached, clear_cache_hook, pop_from_cache_ignore_args
14 from ..utils.warnings import OldVersionWarning
15 from ._variational_strategy import _VariationalStrategy
16 
17 
18 def _ensure_updated_strategy_flag_set(
19     state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
20 ):
21     device = state_dict[list(state_dict.keys())[0]].device
22     if prefix + "updated_strategy" not in state_dict:
23         state_dict[prefix + "updated_strategy"] = torch.tensor(False, device=device)
24         warnings.warn(
25             "You have loaded a variational GP model (using `VariationalStrategy`) from a previous version of "
26             "GPyTorch. We have updated the parameters of your model to work with the new version of "
27             "`VariationalStrategy` that uses whitened parameters.\nYour model will work as expected, but we "
28             "recommend that you re-save your model.",
29             OldVersionWarning,
30         )
31 
32 
33 class VariationalStrategy(_VariationalStrategy):
34     r"""
35     The standard variational strategy, as defined by `Hensman et al. (2015)`_.
36     This strategy takes a set of :math:`m \ll n` inducing points :math:`\mathbf Z`
37     and applies an approximate distribution :math:`q( \mathbf u)` over their function values.
38     (Here, we use the common notation :math:`\mathbf u = f(\mathbf Z)`.
39     The approximate function distribution for any abitrary input :math:`\mathbf X` is given by:
40 
41     .. math::
42 
43         q( f(\mathbf X) ) = \int p( f(\mathbf X) \mid \mathbf u) q(\mathbf u) \: d\mathbf u
44 
45     This variational strategy uses "whitening" to accelerate the optimization of the variational
46     parameters. See `Matthews (2017)`_ for more info.
47 
48     :param ~gpytorch.models.ApproximateGP model: Model this strategy is applied to.
49         Typically passed in when the VariationalStrategy is created in the
50         __init__ method of the user defined model.
51     :param torch.Tensor inducing_points: Tensor containing a set of inducing
52         points to use for variational inference.
53     :param ~gpytorch.variational.VariationalDistribution variational_distribution: A
54         VariationalDistribution object that represents the form of the variational distribution :math:`q(\mathbf u)`
55     :param learn_inducing_locations: (Default True): Whether or not
56         the inducing point locations :math:`\mathbf Z` should be learned (i.e. are they
57         parameters of the model).
58     :type learn_inducing_locations: `bool`, optional
59 
60     .. _Hensman et al. (2015):
61         http://proceedings.mlr.press/v38/hensman15.pdf
62     .. _Matthews (2017):
63         https://www.repository.cam.ac.uk/handle/1810/278022
64     """
65 
66     def __init__(self, model, inducing_points, variational_distribution, learn_inducing_locations=True):
67         super().__init__(model, inducing_points, variational_distribution, learn_inducing_locations)
68         self.register_buffer("updated_strategy", torch.tensor(True))
69         self._register_load_state_dict_pre_hook(_ensure_updated_strategy_flag_set)
70 
71     @cached(name="cholesky_factor", ignore_args=True)
72     def _cholesky_factor(self, induc_induc_covar):
73         L = psd_safe_cholesky(delazify(induc_induc_covar).double(), jitter=settings.cholesky_jitter.value())
74         return TriangularLazyTensor(L)
75 
76     @property
77     @cached(name="prior_distribution_memo")
78     def prior_distribution(self):
79         zeros = torch.zeros(
80             self._variational_distribution.shape(),
81             dtype=self._variational_distribution.dtype,
82             device=self._variational_distribution.device,
83         )
84         ones = torch.ones_like(zeros)
85         res = MultivariateNormal(zeros, DiagLazyTensor(ones))
86         return res
87 
88     def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None, **kwargs):
89         # Compute full prior distribution
90         full_inputs = torch.cat([inducing_points, x], dim=-2)
91         full_output = self.model.forward(full_inputs, **kwargs)
92         full_covar = full_output.lazy_covariance_matrix
93 
94         # Covariance terms
95         num_induc = inducing_points.size(-2)
96         test_mean = full_output.mean[..., num_induc:]
97         induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter()
98         induc_data_covar = full_covar[..., :num_induc, num_induc:].evaluate()
99         data_data_covar = full_covar[..., num_induc:, num_induc:]
100 
101         # Compute interpolation terms
102         # K_ZZ^{-1/2} K_ZX
103         # K_ZZ^{-1/2} \mu_Z
104         L = self._cholesky_factor(induc_induc_covar)
105         if L.shape != induc_induc_covar.shape:
106             # Aggressive caching can cause nasty shape incompatibilies when evaluating with different batch shapes
107             # TODO: Use a hook fo this
108             try:
109                 pop_from_cache_ignore_args(self, "cholesky_factor")
110             except CachingError:
111                 pass
112             L = self._cholesky_factor(induc_induc_covar)
113         interp_term = L.inv_matmul(induc_data_covar.double()).to(full_inputs.dtype)
114 
115         # Compute the mean of q(f)
116         # k_XZ K_ZZ^{-1/2} (m - K_ZZ^{-1/2} \mu_Z) + \mu_X
117         predictive_mean = (
118             torch.matmul(
119                 interp_term.transpose(-1, -2), (inducing_values - self.prior_distribution.mean).unsqueeze(-1)
120             ).squeeze(-1)
121             + test_mean
122         )
123 
124         # Compute the covariance of q(f)
125         # K_XX + k_XZ K_ZZ^{-1/2} (S - I) K_ZZ^{-1/2} k_ZX
126         middle_term = self.prior_distribution.lazy_covariance_matrix.mul(-1)
127         if variational_inducing_covar is not None:
128             middle_term = SumLazyTensor(variational_inducing_covar, middle_term)
129 
130         if trace_mode.on():
131             predictive_covar = (
132                 data_data_covar.add_jitter(1e-4).evaluate()
133                 + interp_term.transpose(-1, -2) @ middle_term.evaluate() @ interp_term
134             )
135         else:
136             predictive_covar = SumLazyTensor(
137                 data_data_covar.add_jitter(1e-4),
138                 MatmulLazyTensor(interp_term.transpose(-1, -2), middle_term @ interp_term),
139             )
140 
141         # Return the distribution
142         return MultivariateNormal(predictive_mean, predictive_covar)
143 
144     def __call__(self, x, prior=False, **kwargs):
145         if not self.updated_strategy.item() and not prior:
146             with torch.no_grad():
147                 # Get unwhitened p(u)
148                 prior_function_dist = self(self.inducing_points, prior=True)
149                 prior_mean = prior_function_dist.loc
150                 L = self._cholesky_factor(prior_function_dist.lazy_covariance_matrix.add_jitter())
151 
152                 # Temporarily turn off noise that's added to the mean
153                 orig_mean_init_std = self._variational_distribution.mean_init_std
154                 self._variational_distribution.mean_init_std = 0.0
155 
156                 # Change the variational parameters to be whitened
157                 variational_dist = self.variational_distribution
158                 mean_diff = (variational_dist.loc - prior_mean).unsqueeze(-1).double()
159                 whitened_mean = L.inv_matmul(mean_diff).squeeze(-1).to(variational_dist.loc.dtype)
160                 covar_root = variational_dist.lazy_covariance_matrix.root_decomposition().root.evaluate().double()
161                 whitened_covar = RootLazyTensor(L.inv_matmul(covar_root).to(variational_dist.loc.dtype))
162                 whitened_variational_distribution = variational_dist.__class__(whitened_mean, whitened_covar)
163                 self._variational_distribution.initialize_variational_distribution(whitened_variational_distribution)
164 
165                 # Reset the random noise parameter of the model
166                 self._variational_distribution.mean_init_std = orig_mean_init_std
167 
168                 # Reset the cache
169                 clear_cache_hook(self)
170 
171                 # Mark that we have updated the variational strategy
172                 self.updated_strategy.fill_(True)
173 
174         return super().__call__(x, prior=prior, **kwargs)
175
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/gpytorch/variational/variational_strategy.py b/gpytorch/variational/variational_strategy.py
--- a/gpytorch/variational/variational_strategy.py
+++ b/gpytorch/variational/variational_strategy.py
@@ -114,12 +114,7 @@
 
         # Compute the mean of q(f)
         # k_XZ K_ZZ^{-1/2} (m - K_ZZ^{-1/2} \mu_Z) + \mu_X
-        predictive_mean = (
-            torch.matmul(
-                interp_term.transpose(-1, -2), (inducing_values - self.prior_distribution.mean).unsqueeze(-1)
-            ).squeeze(-1)
-            + test_mean
-        )
+        predictive_mean = (interp_term.transpose(-1, -2) @ inducing_values.unsqueeze(-1)).squeeze(-1) + test_mean
 
         # Compute the covariance of q(f)
         # K_XX + k_XZ K_ZZ^{-1/2} (S - I) K_ZZ^{-1/2} k_ZX

{"golden_diff": "diff --git a/gpytorch/variational/variational_strategy.py b/gpytorch/variational/variational_strategy.py\n--- a/gpytorch/variational/variational_strategy.py\n+++ b/gpytorch/variational/variational_strategy.py\n@@ -114,12 +114,7 @@\n \n         # Compute the mean of q(f)\n         # k_XZ K_ZZ^{-1/2} (m - K_ZZ^{-1/2} \\mu_Z) + \\mu_X\n-        predictive_mean = (\n-            torch.matmul(\n-                interp_term.transpose(-1, -2), (inducing_values - self.prior_distribution.mean).unsqueeze(-1)\n-            ).squeeze(-1)\n-            + test_mean\n-        )\n+        predictive_mean = (interp_term.transpose(-1, -2) @ inducing_values.unsqueeze(-1)).squeeze(-1) + test_mean\n \n         # Compute the covariance of q(f)\n         # K_XX + k_XZ K_ZZ^{-1/2} (S - I) K_ZZ^{-1/2} k_ZX\n", "issue": "Code in variational_strategy.py\n# \ud83d\udc1b Bug\r\n\r\nDear Gpytorch Developers\r\n\r\nHere is a possible issue with variational_strategy.py\r\n\r\nI am wondering whether it is a special consideration or something missing.  According to the formula in Line 121, it seems an L is missing in front of self.prior_distribution.mean in Line 115.   I understand variable inducing_values has already included L.inv() according to whitening of the inducing posterior, but this is not the case for the prior.   It is fortunate that self.prior_distribution.mean is always 0, so no errors.\r\n\r\nThank you for the great package\r\n\r\nJ.\r\n\r\n## To reproduce\r\n\r\n** Code snippet to reproduce **\r\n```python\r\n# Your code goes here\r\n# Please make sure it does not require any external dependencies (other than PyTorch!)\r\n# (We much prefer small snippets rather than links to existing libraries!)\r\n```\r\n\r\n** Stack trace/error message **\r\n```\r\n// Paste the bad output here!\r\n```\r\n\r\n## Expected Behavior\r\n\r\n<!-- A clear and concise description of what you expected to happen. -->\r\n\r\n## System information\r\n\r\n**Please complete the following information:**\r\n- <!-- GPyTorch Version (run `print(gpytorch.__version__)` -->\r\n- <!-- PyTorch Version (run `print(torch.__version__)` -->\r\n- <!-- Computer OS -->\r\n\r\n## Additional context\r\nAdd any other context about the problem here.\r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n\nimport warnings\n\nimport torch\n\nfrom .. import settings\nfrom ..distributions import MultivariateNormal\nfrom ..lazy import DiagLazyTensor, MatmulLazyTensor, RootLazyTensor, SumLazyTensor, TriangularLazyTensor, delazify\nfrom ..settings import trace_mode\nfrom ..utils.cholesky import psd_safe_cholesky\nfrom ..utils.errors import CachingError\nfrom ..utils.memoize import cached, clear_cache_hook, pop_from_cache_ignore_args\nfrom ..utils.warnings import OldVersionWarning\nfrom ._variational_strategy import _VariationalStrategy\n\n\ndef _ensure_updated_strategy_flag_set(\n    state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs\n):\n    device = state_dict[list(state_dict.keys())[0]].device\n    if prefix + \"updated_strategy\" not in state_dict:\n        state_dict[prefix + \"updated_strategy\"] = torch.tensor(False, device=device)\n        warnings.warn(\n            \"You have loaded a variational GP model (using `VariationalStrategy`) from a previous version of \"\n            \"GPyTorch. We have updated the parameters of your model to work with the new version of \"\n            \"`VariationalStrategy` that uses whitened parameters.\\nYour model will work as expected, but we \"\n            \"recommend that you re-save your model.\",\n            OldVersionWarning,\n        )\n\n\nclass VariationalStrategy(_VariationalStrategy):\n    r\"\"\"\n    The standard variational strategy, as defined by `Hensman et al. (2015)`_.\n    This strategy takes a set of :math:`m \\ll n` inducing points :math:`\\mathbf Z`\n    and applies an approximate distribution :math:`q( \\mathbf u)` over their function values.\n    (Here, we use the common notation :math:`\\mathbf u = f(\\mathbf Z)`.\n    The approximate function distribution for any abitrary input :math:`\\mathbf X` is given by:\n\n    .. math::\n\n        q( f(\\mathbf X) ) = \\int p( f(\\mathbf X) \\mid \\mathbf u) q(\\mathbf u) \\: d\\mathbf u\n\n    This variational strategy uses \"whitening\" to accelerate the optimization of the variational\n    parameters. See `Matthews (2017)`_ for more info.\n\n    :param ~gpytorch.models.ApproximateGP model: Model this strategy is applied to.\n        Typically passed in when the VariationalStrategy is created in the\n        __init__ method of the user defined model.\n    :param torch.Tensor inducing_points: Tensor containing a set of inducing\n        points to use for variational inference.\n    :param ~gpytorch.variational.VariationalDistribution variational_distribution: A\n        VariationalDistribution object that represents the form of the variational distribution :math:`q(\\mathbf u)`\n    :param learn_inducing_locations: (Default True): Whether or not\n        the inducing point locations :math:`\\mathbf Z` should be learned (i.e. are they\n        parameters of the model).\n    :type learn_inducing_locations: `bool`, optional\n\n    .. _Hensman et al. (2015):\n        http://proceedings.mlr.press/v38/hensman15.pdf\n    .. _Matthews (2017):\n        https://www.repository.cam.ac.uk/handle/1810/278022\n    \"\"\"\n\n    def __init__(self, model, inducing_points, variational_distribution, learn_inducing_locations=True):\n        super().__init__(model, inducing_points, variational_distribution, learn_inducing_locations)\n        self.register_buffer(\"updated_strategy\", torch.tensor(True))\n        self._register_load_state_dict_pre_hook(_ensure_updated_strategy_flag_set)\n\n    @cached(name=\"cholesky_factor\", ignore_args=True)\n    def _cholesky_factor(self, induc_induc_covar):\n        L = psd_safe_cholesky(delazify(induc_induc_covar).double(), jitter=settings.cholesky_jitter.value())\n        return TriangularLazyTensor(L)\n\n    @property\n    @cached(name=\"prior_distribution_memo\")\n    def prior_distribution(self):\n        zeros = torch.zeros(\n            self._variational_distribution.shape(),\n            dtype=self._variational_distribution.dtype,\n            device=self._variational_distribution.device,\n        )\n        ones = torch.ones_like(zeros)\n        res = MultivariateNormal(zeros, DiagLazyTensor(ones))\n        return res\n\n    def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None, **kwargs):\n        # Compute full prior distribution\n        full_inputs = torch.cat([inducing_points, x], dim=-2)\n        full_output = self.model.forward(full_inputs, **kwargs)\n        full_covar = full_output.lazy_covariance_matrix\n\n        # Covariance terms\n        num_induc = inducing_points.size(-2)\n        test_mean = full_output.mean[..., num_induc:]\n        induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter()\n        induc_data_covar = full_covar[..., :num_induc, num_induc:].evaluate()\n        data_data_covar = full_covar[..., num_induc:, num_induc:]\n\n        # Compute interpolation terms\n        # K_ZZ^{-1/2} K_ZX\n        # K_ZZ^{-1/2} \\mu_Z\n        L = self._cholesky_factor(induc_induc_covar)\n        if L.shape != induc_induc_covar.shape:\n            # Aggressive caching can cause nasty shape incompatibilies when evaluating with different batch shapes\n            # TODO: Use a hook fo this\n            try:\n                pop_from_cache_ignore_args(self, \"cholesky_factor\")\n            except CachingError:\n                pass\n            L = self._cholesky_factor(induc_induc_covar)\n        interp_term = L.inv_matmul(induc_data_covar.double()).to(full_inputs.dtype)\n\n        # Compute the mean of q(f)\n        # k_XZ K_ZZ^{-1/2} (m - K_ZZ^{-1/2} \\mu_Z) + \\mu_X\n        predictive_mean = (\n            torch.matmul(\n                interp_term.transpose(-1, -2), (inducing_values - self.prior_distribution.mean).unsqueeze(-1)\n            ).squeeze(-1)\n            + test_mean\n        )\n\n        # Compute the covariance of q(f)\n        # K_XX + k_XZ K_ZZ^{-1/2} (S - I) K_ZZ^{-1/2} k_ZX\n        middle_term = self.prior_distribution.lazy_covariance_matrix.mul(-1)\n        if variational_inducing_covar is not None:\n            middle_term = SumLazyTensor(variational_inducing_covar, middle_term)\n\n        if trace_mode.on():\n            predictive_covar = (\n                data_data_covar.add_jitter(1e-4).evaluate()\n                + interp_term.transpose(-1, -2) @ middle_term.evaluate() @ interp_term\n            )\n        else:\n            predictive_covar = SumLazyTensor(\n                data_data_covar.add_jitter(1e-4),\n                MatmulLazyTensor(interp_term.transpose(-1, -2), middle_term @ interp_term),\n            )\n\n        # Return the distribution\n        return MultivariateNormal(predictive_mean, predictive_covar)\n\n    def __call__(self, x, prior=False, **kwargs):\n        if not self.updated_strategy.item() and not prior:\n            with torch.no_grad():\n                # Get unwhitened p(u)\n                prior_function_dist = self(self.inducing_points, prior=True)\n                prior_mean = prior_function_dist.loc\n                L = self._cholesky_factor(prior_function_dist.lazy_covariance_matrix.add_jitter())\n\n                # Temporarily turn off noise that's added to the mean\n                orig_mean_init_std = self._variational_distribution.mean_init_std\n                self._variational_distribution.mean_init_std = 0.0\n\n                # Change the variational parameters to be whitened\n                variational_dist = self.variational_distribution\n                mean_diff = (variational_dist.loc - prior_mean).unsqueeze(-1).double()\n                whitened_mean = L.inv_matmul(mean_diff).squeeze(-1).to(variational_dist.loc.dtype)\n                covar_root = variational_dist.lazy_covariance_matrix.root_decomposition().root.evaluate().double()\n                whitened_covar = RootLazyTensor(L.inv_matmul(covar_root).to(variational_dist.loc.dtype))\n                whitened_variational_distribution = variational_dist.__class__(whitened_mean, whitened_covar)\n                self._variational_distribution.initialize_variational_distribution(whitened_variational_distribution)\n\n                # Reset the random noise parameter of the model\n                self._variational_distribution.mean_init_std = orig_mean_init_std\n\n                # Reset the cache\n                clear_cache_hook(self)\n\n                # Mark that we have updated the variational strategy\n                self.updated_strategy.fill_(True)\n\n        return super().__call__(x, prior=prior, **kwargs)\n", "path": "gpytorch/variational/variational_strategy.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\nimport warnings\n\nimport torch\n\nfrom .. import settings\nfrom ..distributions import MultivariateNormal\nfrom ..lazy import DiagLazyTensor, MatmulLazyTensor, RootLazyTensor, SumLazyTensor, TriangularLazyTensor, delazify\nfrom ..settings import trace_mode\nfrom ..utils.cholesky import psd_safe_cholesky\nfrom ..utils.errors import CachingError\nfrom ..utils.memoize import cached, clear_cache_hook, pop_from_cache_ignore_args\nfrom ..utils.warnings import OldVersionWarning\nfrom ._variational_strategy import _VariationalStrategy\n\n\ndef _ensure_updated_strategy_flag_set(\n    state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs\n):\n    device = state_dict[list(state_dict.keys())[0]].device\n    if prefix + \"updated_strategy\" not in state_dict:\n        state_dict[prefix + \"updated_strategy\"] = torch.tensor(False, device=device)\n        warnings.warn(\n            \"You have loaded a variational GP model (using `VariationalStrategy`) from a previous version of \"\n            \"GPyTorch. We have updated the parameters of your model to work with the new version of \"\n            \"`VariationalStrategy` that uses whitened parameters.\\nYour model will work as expected, but we \"\n            \"recommend that you re-save your model.\",\n            OldVersionWarning,\n        )\n\n\nclass VariationalStrategy(_VariationalStrategy):\n    r\"\"\"\n    The standard variational strategy, as defined by `Hensman et al. (2015)`_.\n    This strategy takes a set of :math:`m \\ll n` inducing points :math:`\\mathbf Z`\n    and applies an approximate distribution :math:`q( \\mathbf u)` over their function values.\n    (Here, we use the common notation :math:`\\mathbf u = f(\\mathbf Z)`.\n    The approximate function distribution for any abitrary input :math:`\\mathbf X` is given by:\n\n    .. math::\n\n        q( f(\\mathbf X) ) = \\int p( f(\\mathbf X) \\mid \\mathbf u) q(\\mathbf u) \\: d\\mathbf u\n\n    This variational strategy uses \"whitening\" to accelerate the optimization of the variational\n    parameters. See `Matthews (2017)`_ for more info.\n\n    :param ~gpytorch.models.ApproximateGP model: Model this strategy is applied to.\n        Typically passed in when the VariationalStrategy is created in the\n        __init__ method of the user defined model.\n    :param torch.Tensor inducing_points: Tensor containing a set of inducing\n        points to use for variational inference.\n    :param ~gpytorch.variational.VariationalDistribution variational_distribution: A\n        VariationalDistribution object that represents the form of the variational distribution :math:`q(\\mathbf u)`\n    :param learn_inducing_locations: (Default True): Whether or not\n        the inducing point locations :math:`\\mathbf Z` should be learned (i.e. are they\n        parameters of the model).\n    :type learn_inducing_locations: `bool`, optional\n\n    .. _Hensman et al. (2015):\n        http://proceedings.mlr.press/v38/hensman15.pdf\n    .. _Matthews (2017):\n        https://www.repository.cam.ac.uk/handle/1810/278022\n    \"\"\"\n\n    def __init__(self, model, inducing_points, variational_distribution, learn_inducing_locations=True):\n        super().__init__(model, inducing_points, variational_distribution, learn_inducing_locations)\n        self.register_buffer(\"updated_strategy\", torch.tensor(True))\n        self._register_load_state_dict_pre_hook(_ensure_updated_strategy_flag_set)\n\n    @cached(name=\"cholesky_factor\", ignore_args=True)\n    def _cholesky_factor(self, induc_induc_covar):\n        L = psd_safe_cholesky(delazify(induc_induc_covar).double(), jitter=settings.cholesky_jitter.value())\n        return TriangularLazyTensor(L)\n\n    @property\n    @cached(name=\"prior_distribution_memo\")\n    def prior_distribution(self):\n        zeros = torch.zeros(\n            self._variational_distribution.shape(),\n            dtype=self._variational_distribution.dtype,\n            device=self._variational_distribution.device,\n        )\n        ones = torch.ones_like(zeros)\n        res = MultivariateNormal(zeros, DiagLazyTensor(ones))\n        return res\n\n    def forward(self, x, inducing_points, inducing_values, variational_inducing_covar=None, **kwargs):\n        # Compute full prior distribution\n        full_inputs = torch.cat([inducing_points, x], dim=-2)\n        full_output = self.model.forward(full_inputs, **kwargs)\n        full_covar = full_output.lazy_covariance_matrix\n\n        # Covariance terms\n        num_induc = inducing_points.size(-2)\n        test_mean = full_output.mean[..., num_induc:]\n        induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter()\n        induc_data_covar = full_covar[..., :num_induc, num_induc:].evaluate()\n        data_data_covar = full_covar[..., num_induc:, num_induc:]\n\n        # Compute interpolation terms\n        # K_ZZ^{-1/2} K_ZX\n        # K_ZZ^{-1/2} \\mu_Z\n        L = self._cholesky_factor(induc_induc_covar)\n        if L.shape != induc_induc_covar.shape:\n            # Aggressive caching can cause nasty shape incompatibilies when evaluating with different batch shapes\n            # TODO: Use a hook fo this\n            try:\n                pop_from_cache_ignore_args(self, \"cholesky_factor\")\n            except CachingError:\n                pass\n            L = self._cholesky_factor(induc_induc_covar)\n        interp_term = L.inv_matmul(induc_data_covar.double()).to(full_inputs.dtype)\n\n        # Compute the mean of q(f)\n        # k_XZ K_ZZ^{-1/2} (m - K_ZZ^{-1/2} \\mu_Z) + \\mu_X\n        predictive_mean = (interp_term.transpose(-1, -2) @ inducing_values.unsqueeze(-1)).squeeze(-1) + test_mean\n\n        # Compute the covariance of q(f)\n        # K_XX + k_XZ K_ZZ^{-1/2} (S - I) K_ZZ^{-1/2} k_ZX\n        middle_term = self.prior_distribution.lazy_covariance_matrix.mul(-1)\n        if variational_inducing_covar is not None:\n            middle_term = SumLazyTensor(variational_inducing_covar, middle_term)\n\n        if trace_mode.on():\n            predictive_covar = (\n                data_data_covar.add_jitter(1e-4).evaluate()\n                + interp_term.transpose(-1, -2) @ middle_term.evaluate() @ interp_term\n            )\n        else:\n            predictive_covar = SumLazyTensor(\n                data_data_covar.add_jitter(1e-4),\n                MatmulLazyTensor(interp_term.transpose(-1, -2), middle_term @ interp_term),\n            )\n\n        # Return the distribution\n        return MultivariateNormal(predictive_mean, predictive_covar)\n\n    def __call__(self, x, prior=False, **kwargs):\n        if not self.updated_strategy.item() and not prior:\n            with torch.no_grad():\n                # Get unwhitened p(u)\n                prior_function_dist = self(self.inducing_points, prior=True)\n                prior_mean = prior_function_dist.loc\n                L = self._cholesky_factor(prior_function_dist.lazy_covariance_matrix.add_jitter())\n\n                # Temporarily turn off noise that's added to the mean\n                orig_mean_init_std = self._variational_distribution.mean_init_std\n                self._variational_distribution.mean_init_std = 0.0\n\n                # Change the variational parameters to be whitened\n                variational_dist = self.variational_distribution\n                mean_diff = (variational_dist.loc - prior_mean).unsqueeze(-1).double()\n                whitened_mean = L.inv_matmul(mean_diff).squeeze(-1).to(variational_dist.loc.dtype)\n                covar_root = variational_dist.lazy_covariance_matrix.root_decomposition().root.evaluate().double()\n                whitened_covar = RootLazyTensor(L.inv_matmul(covar_root).to(variational_dist.loc.dtype))\n                whitened_variational_distribution = variational_dist.__class__(whitened_mean, whitened_covar)\n                self._variational_distribution.initialize_variational_distribution(whitened_variational_distribution)\n\n                # Reset the random noise parameter of the model\n                self._variational_distribution.mean_init_std = orig_mean_init_std\n\n                # Reset the cache\n                clear_cache_hook(self)\n\n                # Mark that we have updated the variational strategy\n                self.updated_strategy.fill_(True)\n\n        return super().__call__(x, prior=prior, **kwargs)\n", "path": "gpytorch/variational/variational_strategy.py"}]}