We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Problem with batch Hadamard Multitask GP regression (likely a bug in GPyTorch)
Hi I was trying to build a batch hadamard multitask GP regression model following this:
https://gpytorch.readthedocs.io/en/latest/examples/03_Multitask_Exact_GPs/Hadamard_Multitask_GP_Regression.html 

Here is my code: 

    import math
    import torch
    import gpytorch
    from matplotlib import pyplot as plt

    # %%

    train_x1 = torch.rand(50)
    train_x2 = torch.rand(10)

    train_y1 = torch.sin(train_x1 * (2 * math.pi)) + \
        torch.randn(train_x1.size()) * 0.2
    train_y2 = torch.cos(train_x2 * (2 * math.pi)) + \
        torch.randn(train_x2.size()) * 0.2

    # %%


    class MultitaskGPModel(gpytorch.models.ExactGP):
        def __init__(self, train_x, train_y, likelihood):
            super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)
            self.mean_module = gpytorch.means.ConstantMean(
                batch_shape=torch.Size([4]))
            self.covar_module = gpytorch.kernels.RBFKernel(
                batch_shape=torch.Size([4]))

            self.task_covar_module = gpytorch.kernels.IndexKernel(
                num_tasks=2, rank=1, batch_shape=torch.Size([4]))

        def forward(self, x, i):
            mean_x = self.mean_module(x)

            # Get input-input covariance
            covar_x = self.covar_module(x)
            # Get task-task covariance
            covar_i = self.task_covar_module(i)
            # Multiply the two together to get the covariance we want
            covar = covar_x.mul(covar_i)

            return gpytorch.distributions.MultivariateNormal(mean_x, covar)


    likelihood = gpytorch.likelihoods.GaussianLikelihood()

    # %% Reshaping to simulate batches
    train_i_task1 = torch.full_like(train_x1, dtype=torch.long, fill_value=0)
    train_i_task2 = torch.full_like(train_x2, dtype=torch.long, fill_value=1)

    full_train_x = torch.cat([train_x1, train_x2])
    full_train_x = full_train_x.view(1, -1, 1).repeat(4, 1, 1)

    full_train_i = torch.cat([train_i_task1, train_i_task2])
    full_train_i = full_train_i.view(1, -1, 1).repeat(4, 1, 1)

    full_train_y = torch.cat([train_y1, train_y2])
    full_train_y = full_train_y.unsqueeze(0)
    full_train_y = torch.cat(
        (full_train_y, full_train_y, full_train_y, full_train_y))
        
    # Here we have two items that we're passing in as train_inputs
    model = MultitaskGPModel((full_train_x, full_train_i),
                             full_train_y, likelihood)

    # %%

    training_iterations = 50

    # Find optimal model hyperparameters
    model.train()
    likelihood.train()

    # Use the adam optimizer
    optimizer = torch.optim.Adam([
        {'params': model.parameters()},  # Includes GaussianLikelihood parameters
    ], lr=0.1)

    # "Loss" for GPs - the marginal log likelihood
    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

    for i in range(training_iterations):
        optimizer.zero_grad()
        output = model(full_train_x, full_train_i)
        loss = -mll(output, full_train_y).sum()
        loss.backward()
        print('Iter %d/50 - Loss: %.3f' % (i + 1, loss.item()))
        optimizer.step()

I just replicated the existing train x, i and y values 4 times to get a batch size of 4. 

However when I run this, it spits out an error: 

    ---------------------------------------------------------------------------
    RuntimeError                              Traceback (most recent call last)
    ~/test_multitask_gp_unequal_samples.py in 
         87 for i in range(training_iterations):
         88     optimizer.zero_grad()
    ---> 89     output = model(full_train_x, full_train_i)
         90     loss = -mll(output, full_train_y).sum()
         91     loss.backward()

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)
        263                 if not all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
        264                     raise RuntimeError("You must train on the training inputs!")
    --> 265             res = super().__call__(*inputs, **kwargs)
        266             return res
        267 

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)
         22 
         23     def __call__(self, *inputs, **kwargs):
    ---> 24         outputs = self.forward(*inputs, **kwargs)
         25         if isinstance(outputs, list):
         26             return [_validate_module_outputs(output) for output in outputs]

    ~/forecast_research/general_experiments/probabilistic_programs/test_multitask_gp_unequal_samples.py in forward(self, x, i)
         44         covar_i = self.task_covar_module(i)
         45         # Multiply the two together to get the covariance we want
    ---> 46         covar = covar_x.mul(covar_i)
         47 
         48         return gpytorch.distributions.MultivariateNormal(mean_x, covar)

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in mul(self, other)
       1145                 return self._mul_constant(other.view(*other.shape[:-2]))
       1146 
    -> 1147         return self._mul_matrix(lazify(other))
       1148 
       1149     def ndimension(self):

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in _mul_matrix(self, other)
        502 
        503         self = self.evaluate_kernel()
    --> 504         other = other.evaluate_kernel()
        505         if isinstance(self, NonLazyTensor) or isinstance(other, NonLazyTensor):
        506             return NonLazyTensor(self.evaluate() * other.evaluate())

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)
         32         cache_name = name if name is not None else method
         33         if not is_in_cache(self, cache_name):
    ---> 34             add_to_cache(self, cache_name, method(self, *args, **kwargs))
         35         return get_from_cache(self, cache_name)
         36 

    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)
        283             if res.shape != self.shape:
        284                 raise RuntimeError(
    --> 285                     f"The expected shape of the kernel was {self.shape}, but got {res.shape}. "
        286                     "This is likely a bug in GPyTorch."
        287                 )

    RuntimeError: The expected shape of the kernel was torch.Size([4, 60, 60]), but got torch.Size([4, 4, 60, 60]). This is likely a bug in GPyTorch. 

I think there is problem when it tries to multiply the two covariance matrices. 

Please let me know what the possible solution could be for this. 

I am using GPyTorch v1.0.0 and PyTorch v1.4


Thanks 


--- 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/kernels/index_kernel.py`
Content:
```
1 #!/usr/bin/env python3
2 
3 import torch
4 
5 from ..constraints import Positive
6 from ..lazy import DiagLazyTensor, InterpolatedLazyTensor, PsdSumLazyTensor, RootLazyTensor
7 from .kernel import Kernel
8 
9 
10 class IndexKernel(Kernel):
11     r"""
12     A kernel for discrete indices. Kernel is defined by a lookup table.
13 
14     .. math::
15 
16         \begin{equation}
17             k(i, j) = \left(BB^\top + \text{diag}(\mathbf v) \right)_{i, j}
18         \end{equation}
19 
20     where :math:`B` is a low-rank matrix, and :math:`\mathbf v` is a  non-negative vector.
21     These parameters are learned.
22 
23     Args:
24         :attr:`num_tasks` (int):
25             Total number of indices.
26         :attr:`batch_shape` (torch.Size, optional):
27             Set if the MultitaskKernel is operating on batches of data (and you want different
28             parameters for each batch)
29         :attr:`rank` (int):
30             Rank of :math:`B` matrix.
31         :attr:`prior` (:obj:`gpytorch.priors.Prior`):
32             Prior for :math:`B` matrix.
33         :attr:`var_constraint` (Constraint, optional):
34             Constraint for added diagonal component. Default: `Positive`.
35 
36     Attributes:
37         covar_factor:
38             The :math:`B` matrix.
39         lov_var:
40             The element-wise log of the :math:`\mathbf v` vector.
41     """
42 
43     def __init__(self, num_tasks, rank=1, prior=None, var_constraint=None, **kwargs):
44         if rank > num_tasks:
45             raise RuntimeError("Cannot create a task covariance matrix larger than the number of tasks")
46         super().__init__(**kwargs)
47 
48         if var_constraint is None:
49             var_constraint = Positive()
50 
51         self.register_parameter(
52             name="covar_factor", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks, rank))
53         )
54         self.register_parameter(name="raw_var", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks)))
55         if prior is not None:
56             self.register_prior("IndexKernelPrior", prior, self._eval_covar_matrix)
57 
58         self.register_constraint("raw_var", var_constraint)
59 
60     @property
61     def var(self):
62         return self.raw_var_constraint.transform(self.raw_var)
63 
64     @var.setter
65     def var(self, value):
66         self._set_var(value)
67 
68     def _set_var(self, value):
69         self.initialize(raw_var=self.raw_var_constraint.inverse_transform(value))
70 
71     def _eval_covar_matrix(self):
72         cf = self.covar_factor
73         return cf @ cf.transpose(-1, -2) + torch.diag_embed(self.var)
74 
75     @property
76     def covar_matrix(self):
77         var = self.var
78         res = PsdSumLazyTensor(RootLazyTensor(self.covar_factor), DiagLazyTensor(var))
79         return res
80 
81     def forward(self, i1, i2, **params):
82         covar_matrix = self._eval_covar_matrix()
83         res = InterpolatedLazyTensor(
84             base_lazy_tensor=covar_matrix,
85             left_interp_indices=i1.expand(self.batch_shape + i1.shape),
86             right_interp_indices=i2.expand(self.batch_shape + i2.shape),
87         )
88         return res
89
```
--- 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/kernels/index_kernel.py b/gpytorch/kernels/index_kernel.py
--- a/gpytorch/kernels/index_kernel.py
+++ b/gpytorch/kernels/index_kernel.py
@@ -80,9 +80,10 @@
 
     def forward(self, i1, i2, **params):
         covar_matrix = self._eval_covar_matrix()
+
         res = InterpolatedLazyTensor(
             base_lazy_tensor=covar_matrix,
-            left_interp_indices=i1.expand(self.batch_shape + i1.shape),
-            right_interp_indices=i2.expand(self.batch_shape + i2.shape),
+            left_interp_indices=i1.expand(self.batch_shape + i1.shape[-2:]),
+            right_interp_indices=i2.expand(self.batch_shape + i2.shape[-2:]),
         )
         return res

{"golden_diff": "diff --git a/gpytorch/kernels/index_kernel.py b/gpytorch/kernels/index_kernel.py\n--- a/gpytorch/kernels/index_kernel.py\n+++ b/gpytorch/kernels/index_kernel.py\n@@ -80,9 +80,10 @@\n \n     def forward(self, i1, i2, **params):\n         covar_matrix = self._eval_covar_matrix()\n+\n         res = InterpolatedLazyTensor(\n             base_lazy_tensor=covar_matrix,\n-            left_interp_indices=i1.expand(self.batch_shape + i1.shape),\n-            right_interp_indices=i2.expand(self.batch_shape + i2.shape),\n+            left_interp_indices=i1.expand(self.batch_shape + i1.shape[-2:]),\n+            right_interp_indices=i2.expand(self.batch_shape + i2.shape[-2:]),\n         )\n         return res\n", "issue": "Problem with batch Hadamard Multitask GP regression (likely a bug in GPyTorch)\nHi I was trying to build a batch hadamard multitask GP regression model following this:\r\nhttps://gpytorch.readthedocs.io/en/latest/examples/03_Multitask_Exact_GPs/Hadamard_Multitask_GP_Regression.html \r\n\r\nHere is my code: \r\n\r\n    import math\r\n    import torch\r\n    import gpytorch\r\n    from matplotlib import pyplot as plt\r\n\r\n    # %%\r\n\r\n    train_x1 = torch.rand(50)\r\n    train_x2 = torch.rand(10)\r\n\r\n    train_y1 = torch.sin(train_x1 * (2 * math.pi)) + \\\r\n        torch.randn(train_x1.size()) * 0.2\r\n    train_y2 = torch.cos(train_x2 * (2 * math.pi)) + \\\r\n        torch.randn(train_x2.size()) * 0.2\r\n\r\n    # %%\r\n\r\n\r\n    class MultitaskGPModel(gpytorch.models.ExactGP):\r\n        def __init__(self, train_x, train_y, likelihood):\r\n            super(MultitaskGPModel, self).__init__(train_x, train_y, likelihood)\r\n            self.mean_module = gpytorch.means.ConstantMean(\r\n                batch_shape=torch.Size([4]))\r\n            self.covar_module = gpytorch.kernels.RBFKernel(\r\n                batch_shape=torch.Size([4]))\r\n\r\n            self.task_covar_module = gpytorch.kernels.IndexKernel(\r\n                num_tasks=2, rank=1, batch_shape=torch.Size([4]))\r\n\r\n        def forward(self, x, i):\r\n            mean_x = self.mean_module(x)\r\n\r\n            # Get input-input covariance\r\n            covar_x = self.covar_module(x)\r\n            # Get task-task covariance\r\n            covar_i = self.task_covar_module(i)\r\n            # Multiply the two together to get the covariance we want\r\n            covar = covar_x.mul(covar_i)\r\n\r\n            return gpytorch.distributions.MultivariateNormal(mean_x, covar)\r\n\r\n\r\n    likelihood = gpytorch.likelihoods.GaussianLikelihood()\r\n\r\n    # %% Reshaping to simulate batches\r\n    train_i_task1 = torch.full_like(train_x1, dtype=torch.long, fill_value=0)\r\n    train_i_task2 = torch.full_like(train_x2, dtype=torch.long, fill_value=1)\r\n\r\n    full_train_x = torch.cat([train_x1, train_x2])\r\n    full_train_x = full_train_x.view(1, -1, 1).repeat(4, 1, 1)\r\n\r\n    full_train_i = torch.cat([train_i_task1, train_i_task2])\r\n    full_train_i = full_train_i.view(1, -1, 1).repeat(4, 1, 1)\r\n\r\n    full_train_y = torch.cat([train_y1, train_y2])\r\n    full_train_y = full_train_y.unsqueeze(0)\r\n    full_train_y = torch.cat(\r\n        (full_train_y, full_train_y, full_train_y, full_train_y))\r\n        \r\n    # Here we have two items that we're passing in as train_inputs\r\n    model = MultitaskGPModel((full_train_x, full_train_i),\r\n                             full_train_y, likelihood)\r\n\r\n    # %%\r\n\r\n    training_iterations = 50\r\n\r\n    # Find optimal model hyperparameters\r\n    model.train()\r\n    likelihood.train()\r\n\r\n    # Use the adam optimizer\r\n    optimizer = torch.optim.Adam([\r\n        {'params': model.parameters()},  # Includes GaussianLikelihood parameters\r\n    ], lr=0.1)\r\n\r\n    # \"Loss\" for GPs - the marginal log likelihood\r\n    mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)\r\n\r\n    for i in range(training_iterations):\r\n        optimizer.zero_grad()\r\n        output = model(full_train_x, full_train_i)\r\n        loss = -mll(output, full_train_y).sum()\r\n        loss.backward()\r\n        print('Iter %d/50 - Loss: %.3f' % (i + 1, loss.item()))\r\n        optimizer.step()\r\n\r\nI just replicated the existing train x, i and y values 4 times to get a batch size of 4. \r\n\r\nHowever when I run this, it spits out an error: \r\n\r\n    ---------------------------------------------------------------------------\r\n    RuntimeError                              Traceback (most recent call last)\r\n    ~/test_multitask_gp_unequal_samples.py in \r\n         87 for i in range(training_iterations):\r\n         88     optimizer.zero_grad()\r\n    ---> 89     output = model(full_train_x, full_train_i)\r\n         90     loss = -mll(output, full_train_y).sum()\r\n         91     loss.backward()\r\n\r\n    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/models/exact_gp.py in __call__(self, *args, **kwargs)\r\n        263                 if not all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):\r\n        264                     raise RuntimeError(\"You must train on the training inputs!\")\r\n    --> 265             res = super().__call__(*inputs, **kwargs)\r\n        266             return res\r\n        267 \r\n\r\n    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/module.py in __call__(self, *inputs, **kwargs)\r\n         22 \r\n         23     def __call__(self, *inputs, **kwargs):\r\n    ---> 24         outputs = self.forward(*inputs, **kwargs)\r\n         25         if isinstance(outputs, list):\r\n         26             return [_validate_module_outputs(output) for output in outputs]\r\n\r\n    ~/forecast_research/general_experiments/probabilistic_programs/test_multitask_gp_unequal_samples.py in forward(self, x, i)\r\n         44         covar_i = self.task_covar_module(i)\r\n         45         # Multiply the two together to get the covariance we want\r\n    ---> 46         covar = covar_x.mul(covar_i)\r\n         47 \r\n         48         return gpytorch.distributions.MultivariateNormal(mean_x, covar)\r\n\r\n    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in mul(self, other)\r\n       1145                 return self._mul_constant(other.view(*other.shape[:-2]))\r\n       1146 \r\n    -> 1147         return self._mul_matrix(lazify(other))\r\n       1148 \r\n       1149     def ndimension(self):\r\n\r\n    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_tensor.py in _mul_matrix(self, other)\r\n        502 \r\n        503         self = self.evaluate_kernel()\r\n    --> 504         other = other.evaluate_kernel()\r\n        505         if isinstance(self, NonLazyTensor) or isinstance(other, NonLazyTensor):\r\n        506             return NonLazyTensor(self.evaluate() * other.evaluate())\r\n\r\n    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/utils/memoize.py in g(self, *args, **kwargs)\r\n         32         cache_name = name if name is not None else method\r\n         33         if not is_in_cache(self, cache_name):\r\n    ---> 34             add_to_cache(self, cache_name, method(self, *args, **kwargs))\r\n         35         return get_from_cache(self, cache_name)\r\n         36 \r\n\r\n    /anaconda3/envs/experimental/lib/python3.7/site-packages/gpytorch/lazy/lazy_evaluated_kernel_tensor.py in evaluate_kernel(self)\r\n        283             if res.shape != self.shape:\r\n        284                 raise RuntimeError(\r\n    --> 285                     f\"The expected shape of the kernel was {self.shape}, but got {res.shape}. \"\r\n        286                     \"This is likely a bug in GPyTorch.\"\r\n        287                 )\r\n\r\n    RuntimeError: The expected shape of the kernel was torch.Size([4, 60, 60]), but got torch.Size([4, 4, 60, 60]). This is likely a bug in GPyTorch. \r\n\r\nI think there is problem when it tries to multiply the two covariance matrices. \r\n\r\nPlease let me know what the possible solution could be for this. \r\n\r\nI am using GPyTorch v1.0.0 and PyTorch v1.4\r\n\r\n\r\nThanks \r\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n\nimport torch\n\nfrom ..constraints import Positive\nfrom ..lazy import DiagLazyTensor, InterpolatedLazyTensor, PsdSumLazyTensor, RootLazyTensor\nfrom .kernel import Kernel\n\n\nclass IndexKernel(Kernel):\n    r\"\"\"\n    A kernel for discrete indices. Kernel is defined by a lookup table.\n\n    .. math::\n\n        \\begin{equation}\n            k(i, j) = \\left(BB^\\top + \\text{diag}(\\mathbf v) \\right)_{i, j}\n        \\end{equation}\n\n    where :math:`B` is a low-rank matrix, and :math:`\\mathbf v` is a  non-negative vector.\n    These parameters are learned.\n\n    Args:\n        :attr:`num_tasks` (int):\n            Total number of indices.\n        :attr:`batch_shape` (torch.Size, optional):\n            Set if the MultitaskKernel is operating on batches of data (and you want different\n            parameters for each batch)\n        :attr:`rank` (int):\n            Rank of :math:`B` matrix.\n        :attr:`prior` (:obj:`gpytorch.priors.Prior`):\n            Prior for :math:`B` matrix.\n        :attr:`var_constraint` (Constraint, optional):\n            Constraint for added diagonal component. Default: `Positive`.\n\n    Attributes:\n        covar_factor:\n            The :math:`B` matrix.\n        lov_var:\n            The element-wise log of the :math:`\\mathbf v` vector.\n    \"\"\"\n\n    def __init__(self, num_tasks, rank=1, prior=None, var_constraint=None, **kwargs):\n        if rank > num_tasks:\n            raise RuntimeError(\"Cannot create a task covariance matrix larger than the number of tasks\")\n        super().__init__(**kwargs)\n\n        if var_constraint is None:\n            var_constraint = Positive()\n\n        self.register_parameter(\n            name=\"covar_factor\", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks, rank))\n        )\n        self.register_parameter(name=\"raw_var\", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks)))\n        if prior is not None:\n            self.register_prior(\"IndexKernelPrior\", prior, self._eval_covar_matrix)\n\n        self.register_constraint(\"raw_var\", var_constraint)\n\n    @property\n    def var(self):\n        return self.raw_var_constraint.transform(self.raw_var)\n\n    @var.setter\n    def var(self, value):\n        self._set_var(value)\n\n    def _set_var(self, value):\n        self.initialize(raw_var=self.raw_var_constraint.inverse_transform(value))\n\n    def _eval_covar_matrix(self):\n        cf = self.covar_factor\n        return cf @ cf.transpose(-1, -2) + torch.diag_embed(self.var)\n\n    @property\n    def covar_matrix(self):\n        var = self.var\n        res = PsdSumLazyTensor(RootLazyTensor(self.covar_factor), DiagLazyTensor(var))\n        return res\n\n    def forward(self, i1, i2, **params):\n        covar_matrix = self._eval_covar_matrix()\n        res = InterpolatedLazyTensor(\n            base_lazy_tensor=covar_matrix,\n            left_interp_indices=i1.expand(self.batch_shape + i1.shape),\n            right_interp_indices=i2.expand(self.batch_shape + i2.shape),\n        )\n        return res\n", "path": "gpytorch/kernels/index_kernel.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n\nimport torch\n\nfrom ..constraints import Positive\nfrom ..lazy import DiagLazyTensor, InterpolatedLazyTensor, PsdSumLazyTensor, RootLazyTensor\nfrom .kernel import Kernel\n\n\nclass IndexKernel(Kernel):\n    r\"\"\"\n    A kernel for discrete indices. Kernel is defined by a lookup table.\n\n    .. math::\n\n        \\begin{equation}\n            k(i, j) = \\left(BB^\\top + \\text{diag}(\\mathbf v) \\right)_{i, j}\n        \\end{equation}\n\n    where :math:`B` is a low-rank matrix, and :math:`\\mathbf v` is a  non-negative vector.\n    These parameters are learned.\n\n    Args:\n        :attr:`num_tasks` (int):\n            Total number of indices.\n        :attr:`batch_shape` (torch.Size, optional):\n            Set if the MultitaskKernel is operating on batches of data (and you want different\n            parameters for each batch)\n        :attr:`rank` (int):\n            Rank of :math:`B` matrix.\n        :attr:`prior` (:obj:`gpytorch.priors.Prior`):\n            Prior for :math:`B` matrix.\n        :attr:`var_constraint` (Constraint, optional):\n            Constraint for added diagonal component. Default: `Positive`.\n\n    Attributes:\n        covar_factor:\n            The :math:`B` matrix.\n        lov_var:\n            The element-wise log of the :math:`\\mathbf v` vector.\n    \"\"\"\n\n    def __init__(self, num_tasks, rank=1, prior=None, var_constraint=None, **kwargs):\n        if rank > num_tasks:\n            raise RuntimeError(\"Cannot create a task covariance matrix larger than the number of tasks\")\n        super().__init__(**kwargs)\n\n        if var_constraint is None:\n            var_constraint = Positive()\n\n        self.register_parameter(\n            name=\"covar_factor\", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks, rank))\n        )\n        self.register_parameter(name=\"raw_var\", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks)))\n        if prior is not None:\n            self.register_prior(\"IndexKernelPrior\", prior, self._eval_covar_matrix)\n\n        self.register_constraint(\"raw_var\", var_constraint)\n\n    @property\n    def var(self):\n        return self.raw_var_constraint.transform(self.raw_var)\n\n    @var.setter\n    def var(self, value):\n        self._set_var(value)\n\n    def _set_var(self, value):\n        self.initialize(raw_var=self.raw_var_constraint.inverse_transform(value))\n\n    def _eval_covar_matrix(self):\n        cf = self.covar_factor\n        return cf @ cf.transpose(-1, -2) + torch.diag_embed(self.var)\n\n    @property\n    def covar_matrix(self):\n        var = self.var\n        res = PsdSumLazyTensor(RootLazyTensor(self.covar_factor), DiagLazyTensor(var))\n        return res\n\n    def forward(self, i1, i2, **params):\n        covar_matrix = self._eval_covar_matrix()\n\n        res = InterpolatedLazyTensor(\n            base_lazy_tensor=covar_matrix,\n            left_interp_indices=i1.expand(self.batch_shape + i1.shape[-2:]),\n            right_interp_indices=i2.expand(self.batch_shape + i2.shape[-2:]),\n        )\n        return res\n", "path": "gpytorch/kernels/index_kernel.py"}]}