We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
ARD for RBFKernelGrad
Hi,
Is it possible to use ARD when training a GP model with derivative data (RBFkernelgrad)? Would it be possible to add support for that?

--- 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/rbf_kernel_grad.py`
Content:
```
1 #!/usr/bin/env python3
2 from .rbf_kernel import RBFKernel
3 import torch
4 from ..lazy.kronecker_product_lazy_tensor import KroneckerProductLazyTensor
5 
6 
7 class RBFKernelGrad(RBFKernel):
8     r"""
9     Computes a covariance matrix of the RBF kernel that models the covariance
10     between the values and partial derivatives for inputs :math:`\mathbf{x_1}`
11     and :math:`\mathbf{x_2}`.
12 
13     See :class:`gpytorch.kernels.Kernel` for descriptions of the lengthscale options.
14 
15     .. note::
16 
17         This kernel does not have an `outputscale` parameter. To add a scaling parameter,
18         decorate this kernel with a :class:`gpytorch.kernels.ScaleKernel`.
19 
20     Args:
21         :attr:`batch_shape` (torch.Size, optional):
22             Set this if you want a separate lengthscale for each
23              batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.
24         :attr:`active_dims` (tuple of ints, optional):
25             Set this if you want to compute the covariance of only a few input dimensions. The ints
26             corresponds to the indices of the dimensions. Default: `None`.
27         :attr:`lengthscale_prior` (Prior, optional):
28             Set this if you want to apply a prior to the lengthscale parameter.  Default: `None`.
29         :attr:`param_transform` (function, optional):
30             Set this if you want to use something other than softplus to ensure positiveness of parameters.
31         :attr:`inv_param_transform` (function, optional):
32             Set this to allow setting parameters directly in transformed space and sampling from priors.
33             Automatically inferred for common transformations such as torch.exp or torch.nn.functional.softplus.
34         :attr:`eps` (float):
35             The minimum value that the lengthscale can take (prevents divide by zero errors). Default: `1e-6`.
36 
37     Attributes:
38         :attr:`lengthscale` (Tensor):
39             The lengthscale parameter. Size/shape of parameter depends on the
40             :attr:`ard_num_dims` and :attr:`batch_shape` arguments.
41 
42     Example:
43         >>> x = torch.randn(10, 5)
44         >>> # Non-batch: Simple option
45         >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad())
46         >>> covar = covar_module(x)  # Output: LazyTensor of size (60 x 60), where 60 = n * (d + 1)
47         >>>
48         >>> batch_x = torch.randn(2, 10, 5)
49         >>> # Batch: Simple option
50         >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad())
51         >>> # Batch: different lengthscale for each batch
52         >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad(batch_shape=torch.Size([2])))
53         >>> covar = covar_module(x)  # Output: LazyTensor of size (2 x 60 x 60)
54     """
55     def forward(self, x1, x2, diag=False, **params):
56         b = 1
57         if len(x1.size()) == 2:
58             n1, d = x1.size()
59             n2, d = x2.size()
60         else:
61             b, n1, d = x1.size()
62             _, n2, _ = x2.size()
63 
64         K = torch.zeros(b, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)  # batch x n1(d+1) x n2(d+1)
65         ell = self.lengthscale.squeeze(-1)
66 
67         if not diag:
68             # Scale the inputs by the lengthscale (for stability)
69             x1_ = x1 / ell
70             x2_ = x2 / ell
71 
72             # Form all possible rank-1 products for the gradient and Hessian blocks
73             outer = x1_.view([b, n1, 1, d]) - x2_.view([b, 1, n2, d])
74             outer = torch.transpose(outer, -1, -2).contiguous()
75 
76             # 1) Kernel block
77             diff = self._covar_dist(x1_, x2_, square_dist=True, **params)
78             K_11 = diff.div_(-2).exp_()
79             K[..., :n1, :n2] = K_11
80 
81             # 2) First gradient block
82             outer1 = outer.view([b, n1, n2 * d]) / ell
83             K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])
84 
85             # 3) Second gradient block
86             outer2 = outer.transpose(-1, -3).contiguous().view([b, n2, n1 * d])
87             outer2 = outer2.transpose(-1, -2) / ell
88             K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])
89 
90             # 4) Hessian block
91             outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])
92             kp = KroneckerProductLazyTensor(
93                 torch.eye(d, d, device=x1.device, dtype=x1.dtype),
94                 torch.ones(n1, n2, device=x1.device, dtype=x1.dtype)
95             )
96             chain_rule = kp.evaluate() / ell.pow(2) - outer3
97             K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])
98 
99             # Symmetrize for stability
100             if n1 == n2 and torch.eq(x1, x2).all():
101                 K = 0.5 * (K.transpose(-1, -2) + K)
102 
103             # Apply a perfect shuffle permutation to match the MutiTask ordering
104             pi1 = torch.arange(n1 * (d + 1)).view(d + 1, n1).t().contiguous().view((n1 * (d + 1)))
105             pi2 = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))
106             K = K[..., pi1, :][..., :, pi2]
107 
108             return K
109 
110         else:  # TODO: This will change when ARD is supported
111             if not (n1 == n2 and torch.eq(x1, x2).all()):
112                 raise RuntimeError("diag=True only works when x1 == x2")
113 
114             kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)
115             grad_diag = torch.ones(1, n2 * d, device=x1.device, dtype=x1.dtype) / (ell.pow(2))
116             k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)
117             pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))
118             return k_diag[..., pi]
119 
120     def size(self, x1, x2):
121         """
122         Given `x_1` with `n_1` data points and `x_2` with `n_2` data points, both in
123         `d` dimensions, RBFKernelGrad returns an `n_1(d+1) x n_2(d+1)` kernel matrix.
124         """
125         non_batch_size = ((x1.size(-1) + 1) * x1.size(-2), (x2.size(-1) + 1) * x2.size(-2))
126         if x1.ndimension() == 3:
127             return torch.Size((x1.size(0),) + non_batch_size)
128         else:
129             return torch.Size(non_batch_size)
130
```
--- 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/rbf_kernel_grad.py b/gpytorch/kernels/rbf_kernel_grad.py
--- a/gpytorch/kernels/rbf_kernel_grad.py
+++ b/gpytorch/kernels/rbf_kernel_grad.py
@@ -62,15 +62,15 @@
             _, n2, _ = x2.size()
 
         K = torch.zeros(b, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)  # batch x n1(d+1) x n2(d+1)
-        ell = self.lengthscale.squeeze(-1)
 
         if not diag:
             # Scale the inputs by the lengthscale (for stability)
-            x1_ = x1 / ell
-            x2_ = x2 / ell
+            x1_ = x1.div(self.lengthscale)
+            x2_ = x2.div(self.lengthscale)
 
             # Form all possible rank-1 products for the gradient and Hessian blocks
             outer = x1_.view([b, n1, 1, d]) - x2_.view([b, 1, n2, d])
+            outer = outer / self.lengthscale
             outer = torch.transpose(outer, -1, -2).contiguous()
 
             # 1) Kernel block
@@ -79,21 +79,21 @@
             K[..., :n1, :n2] = K_11
 
             # 2) First gradient block
-            outer1 = outer.view([b, n1, n2 * d]) / ell
+            outer1 = outer.view([b, n1, n2 * d])
             K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])
 
             # 3) Second gradient block
             outer2 = outer.transpose(-1, -3).contiguous().view([b, n2, n1 * d])
-            outer2 = outer2.transpose(-1, -2) / ell
+            outer2 = outer2.transpose(-1, -2)
             K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])
 
             # 4) Hessian block
             outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])
             kp = KroneckerProductLazyTensor(
-                torch.eye(d, d, device=x1.device, dtype=x1.dtype),
-                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype)
+                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1]) / self.lengthscale.pow_(2),
+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1])
             )
-            chain_rule = kp.evaluate() / ell.pow(2) - outer3
+            chain_rule = kp.evaluate() - outer3
             K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])
 
             # Symmetrize for stability
@@ -107,12 +107,13 @@
 
             return K
 
-        else:  # TODO: This will change when ARD is supported
+        else:
             if not (n1 == n2 and torch.eq(x1, x2).all()):
                 raise RuntimeError("diag=True only works when x1 == x2")
 
             kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)
-            grad_diag = torch.ones(1, n2 * d, device=x1.device, dtype=x1.dtype) / (ell.pow(2))
+            grad_diag = torch.ones(b, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow_(2)
+            grad_diag = grad_diag.transpose(-1, -2).contiguous().view([b, n2 * d])
             k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)
             pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))
             return k_diag[..., pi]

{"golden_diff": "diff --git a/gpytorch/kernels/rbf_kernel_grad.py b/gpytorch/kernels/rbf_kernel_grad.py\n--- a/gpytorch/kernels/rbf_kernel_grad.py\n+++ b/gpytorch/kernels/rbf_kernel_grad.py\n@@ -62,15 +62,15 @@\n             _, n2, _ = x2.size()\n \n         K = torch.zeros(b, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)  # batch x n1(d+1) x n2(d+1)\n-        ell = self.lengthscale.squeeze(-1)\n \n         if not diag:\n             # Scale the inputs by the lengthscale (for stability)\n-            x1_ = x1 / ell\n-            x2_ = x2 / ell\n+            x1_ = x1.div(self.lengthscale)\n+            x2_ = x2.div(self.lengthscale)\n \n             # Form all possible rank-1 products for the gradient and Hessian blocks\n             outer = x1_.view([b, n1, 1, d]) - x2_.view([b, 1, n2, d])\n+            outer = outer / self.lengthscale\n             outer = torch.transpose(outer, -1, -2).contiguous()\n \n             # 1) Kernel block\n@@ -79,21 +79,21 @@\n             K[..., :n1, :n2] = K_11\n \n             # 2) First gradient block\n-            outer1 = outer.view([b, n1, n2 * d]) / ell\n+            outer1 = outer.view([b, n1, n2 * d])\n             K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])\n \n             # 3) Second gradient block\n             outer2 = outer.transpose(-1, -3).contiguous().view([b, n2, n1 * d])\n-            outer2 = outer2.transpose(-1, -2) / ell\n+            outer2 = outer2.transpose(-1, -2)\n             K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])\n \n             # 4) Hessian block\n             outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])\n             kp = KroneckerProductLazyTensor(\n-                torch.eye(d, d, device=x1.device, dtype=x1.dtype),\n-                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype)\n+                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1]) / self.lengthscale.pow_(2),\n+                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1])\n             )\n-            chain_rule = kp.evaluate() / ell.pow(2) - outer3\n+            chain_rule = kp.evaluate() - outer3\n             K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])\n \n             # Symmetrize for stability\n@@ -107,12 +107,13 @@\n \n             return K\n \n-        else:  # TODO: This will change when ARD is supported\n+        else:\n             if not (n1 == n2 and torch.eq(x1, x2).all()):\n                 raise RuntimeError(\"diag=True only works when x1 == x2\")\n \n             kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)\n-            grad_diag = torch.ones(1, n2 * d, device=x1.device, dtype=x1.dtype) / (ell.pow(2))\n+            grad_diag = torch.ones(b, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow_(2)\n+            grad_diag = grad_diag.transpose(-1, -2).contiguous().view([b, n2 * d])\n             k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)\n             pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))\n             return k_diag[..., pi]\n", "issue": "ARD for RBFKernelGrad\nHi,\r\nIs it possible to use ARD when training a GP model with derivative data (RBFkernelgrad)? Would it be possible to add support for that?\n", "before_files": [{"content": "#!/usr/bin/env python3\nfrom .rbf_kernel import RBFKernel\nimport torch\nfrom ..lazy.kronecker_product_lazy_tensor import KroneckerProductLazyTensor\n\n\nclass RBFKernelGrad(RBFKernel):\n    r\"\"\"\n    Computes a covariance matrix of the RBF kernel that models the covariance\n    between the values and partial derivatives for inputs :math:`\\mathbf{x_1}`\n    and :math:`\\mathbf{x_2}`.\n\n    See :class:`gpytorch.kernels.Kernel` for descriptions of the lengthscale options.\n\n    .. note::\n\n        This kernel does not have an `outputscale` parameter. To add a scaling parameter,\n        decorate this kernel with a :class:`gpytorch.kernels.ScaleKernel`.\n\n    Args:\n        :attr:`batch_shape` (torch.Size, optional):\n            Set this if you want a separate lengthscale for each\n             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.\n        :attr:`active_dims` (tuple of ints, optional):\n            Set this if you want to compute the covariance of only a few input dimensions. The ints\n            corresponds to the indices of the dimensions. Default: `None`.\n        :attr:`lengthscale_prior` (Prior, optional):\n            Set this if you want to apply a prior to the lengthscale parameter.  Default: `None`.\n        :attr:`param_transform` (function, optional):\n            Set this if you want to use something other than softplus to ensure positiveness of parameters.\n        :attr:`inv_param_transform` (function, optional):\n            Set this to allow setting parameters directly in transformed space and sampling from priors.\n            Automatically inferred for common transformations such as torch.exp or torch.nn.functional.softplus.\n        :attr:`eps` (float):\n            The minimum value that the lengthscale can take (prevents divide by zero errors). Default: `1e-6`.\n\n    Attributes:\n        :attr:`lengthscale` (Tensor):\n            The lengthscale parameter. Size/shape of parameter depends on the\n            :attr:`ard_num_dims` and :attr:`batch_shape` arguments.\n\n    Example:\n        >>> x = torch.randn(10, 5)\n        >>> # Non-batch: Simple option\n        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad())\n        >>> covar = covar_module(x)  # Output: LazyTensor of size (60 x 60), where 60 = n * (d + 1)\n        >>>\n        >>> batch_x = torch.randn(2, 10, 5)\n        >>> # Batch: Simple option\n        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad())\n        >>> # Batch: different lengthscale for each batch\n        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad(batch_shape=torch.Size([2])))\n        >>> covar = covar_module(x)  # Output: LazyTensor of size (2 x 60 x 60)\n    \"\"\"\n    def forward(self, x1, x2, diag=False, **params):\n        b = 1\n        if len(x1.size()) == 2:\n            n1, d = x1.size()\n            n2, d = x2.size()\n        else:\n            b, n1, d = x1.size()\n            _, n2, _ = x2.size()\n\n        K = torch.zeros(b, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)  # batch x n1(d+1) x n2(d+1)\n        ell = self.lengthscale.squeeze(-1)\n\n        if not diag:\n            # Scale the inputs by the lengthscale (for stability)\n            x1_ = x1 / ell\n            x2_ = x2 / ell\n\n            # Form all possible rank-1 products for the gradient and Hessian blocks\n            outer = x1_.view([b, n1, 1, d]) - x2_.view([b, 1, n2, d])\n            outer = torch.transpose(outer, -1, -2).contiguous()\n\n            # 1) Kernel block\n            diff = self._covar_dist(x1_, x2_, square_dist=True, **params)\n            K_11 = diff.div_(-2).exp_()\n            K[..., :n1, :n2] = K_11\n\n            # 2) First gradient block\n            outer1 = outer.view([b, n1, n2 * d]) / ell\n            K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])\n\n            # 3) Second gradient block\n            outer2 = outer.transpose(-1, -3).contiguous().view([b, n2, n1 * d])\n            outer2 = outer2.transpose(-1, -2) / ell\n            K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])\n\n            # 4) Hessian block\n            outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])\n            kp = KroneckerProductLazyTensor(\n                torch.eye(d, d, device=x1.device, dtype=x1.dtype),\n                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype)\n            )\n            chain_rule = kp.evaluate() / ell.pow(2) - outer3\n            K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])\n\n            # Symmetrize for stability\n            if n1 == n2 and torch.eq(x1, x2).all():\n                K = 0.5 * (K.transpose(-1, -2) + K)\n\n            # Apply a perfect shuffle permutation to match the MutiTask ordering\n            pi1 = torch.arange(n1 * (d + 1)).view(d + 1, n1).t().contiguous().view((n1 * (d + 1)))\n            pi2 = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))\n            K = K[..., pi1, :][..., :, pi2]\n\n            return K\n\n        else:  # TODO: This will change when ARD is supported\n            if not (n1 == n2 and torch.eq(x1, x2).all()):\n                raise RuntimeError(\"diag=True only works when x1 == x2\")\n\n            kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)\n            grad_diag = torch.ones(1, n2 * d, device=x1.device, dtype=x1.dtype) / (ell.pow(2))\n            k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)\n            pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))\n            return k_diag[..., pi]\n\n    def size(self, x1, x2):\n        \"\"\"\n        Given `x_1` with `n_1` data points and `x_2` with `n_2` data points, both in\n        `d` dimensions, RBFKernelGrad returns an `n_1(d+1) x n_2(d+1)` kernel matrix.\n        \"\"\"\n        non_batch_size = ((x1.size(-1) + 1) * x1.size(-2), (x2.size(-1) + 1) * x2.size(-2))\n        if x1.ndimension() == 3:\n            return torch.Size((x1.size(0),) + non_batch_size)\n        else:\n            return torch.Size(non_batch_size)\n", "path": "gpytorch/kernels/rbf_kernel_grad.py"}], "after_files": [{"content": "#!/usr/bin/env python3\nfrom .rbf_kernel import RBFKernel\nimport torch\nfrom ..lazy.kronecker_product_lazy_tensor import KroneckerProductLazyTensor\n\n\nclass RBFKernelGrad(RBFKernel):\n    r\"\"\"\n    Computes a covariance matrix of the RBF kernel that models the covariance\n    between the values and partial derivatives for inputs :math:`\\mathbf{x_1}`\n    and :math:`\\mathbf{x_2}`.\n\n    See :class:`gpytorch.kernels.Kernel` for descriptions of the lengthscale options.\n\n    .. note::\n\n        This kernel does not have an `outputscale` parameter. To add a scaling parameter,\n        decorate this kernel with a :class:`gpytorch.kernels.ScaleKernel`.\n\n    Args:\n        :attr:`batch_shape` (torch.Size, optional):\n            Set this if you want a separate lengthscale for each\n             batch of input data. It should be `b` if :attr:`x1` is a `b x n x d` tensor. Default: `torch.Size([1])`.\n        :attr:`active_dims` (tuple of ints, optional):\n            Set this if you want to compute the covariance of only a few input dimensions. The ints\n            corresponds to the indices of the dimensions. Default: `None`.\n        :attr:`lengthscale_prior` (Prior, optional):\n            Set this if you want to apply a prior to the lengthscale parameter.  Default: `None`.\n        :attr:`param_transform` (function, optional):\n            Set this if you want to use something other than softplus to ensure positiveness of parameters.\n        :attr:`inv_param_transform` (function, optional):\n            Set this to allow setting parameters directly in transformed space and sampling from priors.\n            Automatically inferred for common transformations such as torch.exp or torch.nn.functional.softplus.\n        :attr:`eps` (float):\n            The minimum value that the lengthscale can take (prevents divide by zero errors). Default: `1e-6`.\n\n    Attributes:\n        :attr:`lengthscale` (Tensor):\n            The lengthscale parameter. Size/shape of parameter depends on the\n            :attr:`ard_num_dims` and :attr:`batch_shape` arguments.\n\n    Example:\n        >>> x = torch.randn(10, 5)\n        >>> # Non-batch: Simple option\n        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad())\n        >>> covar = covar_module(x)  # Output: LazyTensor of size (60 x 60), where 60 = n * (d + 1)\n        >>>\n        >>> batch_x = torch.randn(2, 10, 5)\n        >>> # Batch: Simple option\n        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad())\n        >>> # Batch: different lengthscale for each batch\n        >>> covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernelGrad(batch_shape=torch.Size([2])))\n        >>> covar = covar_module(x)  # Output: LazyTensor of size (2 x 60 x 60)\n    \"\"\"\n    def forward(self, x1, x2, diag=False, **params):\n        b = 1\n        if len(x1.size()) == 2:\n            n1, d = x1.size()\n            n2, d = x2.size()\n        else:\n            b, n1, d = x1.size()\n            _, n2, _ = x2.size()\n\n        K = torch.zeros(b, n1 * (d + 1), n2 * (d + 1), device=x1.device, dtype=x1.dtype)  # batch x n1(d+1) x n2(d+1)\n\n        if not diag:\n            # Scale the inputs by the lengthscale (for stability)\n            x1_ = x1.div(self.lengthscale)\n            x2_ = x2.div(self.lengthscale)\n\n            # Form all possible rank-1 products for the gradient and Hessian blocks\n            outer = x1_.view([b, n1, 1, d]) - x2_.view([b, 1, n2, d])\n            outer = outer / self.lengthscale\n            outer = torch.transpose(outer, -1, -2).contiguous()\n\n            # 1) Kernel block\n            diff = self._covar_dist(x1_, x2_, square_dist=True, **params)\n            K_11 = diff.div_(-2).exp_()\n            K[..., :n1, :n2] = K_11\n\n            # 2) First gradient block\n            outer1 = outer.view([b, n1, n2 * d])\n            K[..., :n1, n2:] = outer1 * K_11.repeat([1, 1, d])\n\n            # 3) Second gradient block\n            outer2 = outer.transpose(-1, -3).contiguous().view([b, n2, n1 * d])\n            outer2 = outer2.transpose(-1, -2)\n            K[..., n1:, :n2] = -outer2 * K_11.repeat([1, d, 1])\n\n            # 4) Hessian block\n            outer3 = outer1.repeat([1, d, 1]) * outer2.repeat([1, 1, d])\n            kp = KroneckerProductLazyTensor(\n                torch.eye(d, d, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1]) / self.lengthscale.pow_(2),\n                torch.ones(n1, n2, device=x1.device, dtype=x1.dtype).repeat([b, 1, 1])\n            )\n            chain_rule = kp.evaluate() - outer3\n            K[..., n1:, n2:] = chain_rule * K_11.repeat([1, d, d])\n\n            # Symmetrize for stability\n            if n1 == n2 and torch.eq(x1, x2).all():\n                K = 0.5 * (K.transpose(-1, -2) + K)\n\n            # Apply a perfect shuffle permutation to match the MutiTask ordering\n            pi1 = torch.arange(n1 * (d + 1)).view(d + 1, n1).t().contiguous().view((n1 * (d + 1)))\n            pi2 = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))\n            K = K[..., pi1, :][..., :, pi2]\n\n            return K\n\n        else:\n            if not (n1 == n2 and torch.eq(x1, x2).all()):\n                raise RuntimeError(\"diag=True only works when x1 == x2\")\n\n            kernel_diag = super(RBFKernelGrad, self).forward(x1, x2, diag=True)\n            grad_diag = torch.ones(b, n2, d, device=x1.device, dtype=x1.dtype) / self.lengthscale.pow_(2)\n            grad_diag = grad_diag.transpose(-1, -2).contiguous().view([b, n2 * d])\n            k_diag = torch.cat((kernel_diag, grad_diag), dim=-1)\n            pi = torch.arange(n2 * (d + 1)).view(d + 1, n2).t().contiguous().view((n2 * (d + 1)))\n            return k_diag[..., pi]\n\n    def size(self, x1, x2):\n        \"\"\"\n        Given `x_1` with `n_1` data points and `x_2` with `n_2` data points, both in\n        `d` dimensions, RBFKernelGrad returns an `n_1(d+1) x n_2(d+1)` kernel matrix.\n        \"\"\"\n        non_batch_size = ((x1.size(-1) + 1) * x1.size(-2), (x2.size(-1) + 1) * x2.size(-2))\n        if x1.ndimension() == 3:\n            return torch.Size((x1.size(0),) + non_batch_size)\n        else:\n            return torch.Size(non_batch_size)\n", "path": "gpytorch/kernels/rbf_kernel_grad.py"}]}