Repository: cornellius-gp/gpytorch
Pull request #517 (2019-02-13 23:28:35)

Custom RBF and Matern Functions
This pull request tries to reduce the memory requirements of a RBFKernel and a MaternKernel by introducing a custom backward function when the lengthscales require grad.
- [x] Implement RBFCovariance Function. Fallback to RBFKernel when using ARD or when x1, x2 require grad
- [x] Implement MaternCovariance Function. Fallback to MaternKernel when using ARD or when x1, x2 require grad
- [x] Add unittests for forward and backward

ARD is not yet implemented because I'm not sure if there's a better way to do it than to create a distance matrix for each dimension (resulting in a 3D Tensor) since we need gradients of a 2D Tensor with respect to each of the lengthscales.

Balandat: Haven't really looked into this, just wanted to make a side comment here: In light of #492, we'll want to start support arbitrary batch shapes also in the kernels - so I guess this is a call to be cognizant of that, and to make sure the code will be easily adaptable to these coming changes. 
KeAWang: Fixed. Thanks for pointing that out. Also, I noticed that `MaternKernel.forward` might not actually be arbitrary batch size compatible because how it computes the mean will mess up the broadcasting. I added a commit that fixes that.

I'm worried that there might be a lot of other code that assume there's only one batch dimension but we don't know about them because our tests all assume there's only batch dimension. @jacobrgardner @gpleiss ?
gpleiss: @KeAWang #492 
gpleiss: I'd rather not merge in this PR now - I think we should test against this PR for a while to see if we can make further memory gains.
Balandat: > I'm worried that there might be a lot of other code that assume there's only one batch dimension but we don't know about them because our tests all assume there's only batch dimension.

Yes, I think outside of the LazyTensor refactor most other code assumes a single batch dimension, this will all need to be changed. Most of the hard work should be done in @gpleiss's refactor though.
KeAWang: Here's a quick memory usage benchmark from the output of `torch.cuda.max_memory_allocated` on 10000 1 dimensional training points:

Using RBFKernel on master:
   after forward: 1.2 GB
   after backward: 2.1 GB
Using RBFCovariance:
   after forward: 1.2 GB
   after backward: 1.6 GB
Using MaternKernel (nu=1/2) on master:
   after forward: 2.4 GB
   after backward: 2.8 GB
Using MaternFunction (nu=1/2):
   after forward: 1.2 GB
   after backward: 1.6 GB
Using MaternKernel (nu=3/2) on master:
   after forward: 2.8 GB
   after backward: 3.6 GB
Using MaternFunction (nu=3/2):
   after forward: 1.2 GB
   after backward: 1.6 GB
Using MaternKernel (nu=5/2) on master:
   after forward: 3.2 GB
   after backward: 4.4 GB
Using MaternCovariance (nu=5/2):
   forward: 1.6 GB
   backward: 1.6 GB
jacobrgardner: @KeAWang Can you make sure to keep this branch up to date with any changes you make to the kernels on the `implicit_kernel` branch? It'll probably be cleaner to merge those features separately if possible. If not, we can do that whole branch at once I suppose.
KeAWang: Yep I've been pushing to this remote branch and then merging into the `implicit_kernel` branch
KeAWang: This should be ready for review now
jacobrgardner: This looks good to me other than my comment above about the distance postprocessing functions.
KeAWang: Hm I'm not sure why this was there. It might have been left over from my earlier implementation of RBFCovariance. I reverted this change but now `test_fantasy_updates` mysteriously fails... 
KeAWang: Fixed.

Patch:
diff --git a/README.md b/README.md
index 444bdf11..ba288f34 100644
--- a/README.md
+++ b/README.md
@@ -92,7 +92,7 @@ GPyTorch is primarily maintained by:
 
 <img width="300" src=https://brand.cornell.edu/assets/images/downloads/logos/cornell_logo_simple/cornell_logo_simple.svg alt="Cornell Logo" />
 
-We would like to thank our other contributors including (but not limited to) Eytan Bakshy, David Arbour, Ruihan Wu, Bram Wallace, Sam Stanton, and Jared Frank.
+We would like to thank our other contributors including (but not limited to)  David Arbour, Eytan Bakshy, Jared Frank, Sam Stanton, Bram Wallace, Ke Alexander Wang, Ruihan Wu.
 
 ## Acknowledgements
 Development of GPyTorch is supported by funding from [Facebook](https://research.fb.com/), the [Bill and Melinda Gates Foundation](https://www.gatesfoundation.org/), the [National Science Foundation](https://www.nsf.gov/), and [SAP](https://www.sap.com/index.html).
diff --git a/gpytorch/functions/__init__.py b/gpytorch/functions/__init__.py
index 01023140..9bbaeed9 100644
--- a/gpytorch/functions/__init__.py
+++ b/gpytorch/functions/__init__.py
@@ -6,6 +6,8 @@ from ._dsmm import DSMM
 from ._normal_cdf import NormalCDF
 from ._log_normal_cdf import LogNormalCDF
 from ..utils.deprecation import _deprecated_function_for
+from .rbf_covariance import RBFCovariance
+from .matern_covariance import MaternCovariance
 
 
 def add_diag(input, diag):
@@ -196,6 +198,8 @@ inv_quad_log_det = _deprecated_function_for("inv_quad_log_det", inv_quad_logdet)
 
 
 __all__ = [
+    "MaternCovariance",
+    "RBFCovariance",
     "add_diag",
     "dsmm",
     "inv_matmul",
diff --git a/gpytorch/functions/matern_covariance.py b/gpytorch/functions/matern_covariance.py
new file mode 100644
index 00000000..2e0982e3
--- /dev/null
+++ b/gpytorch/functions/matern_covariance.py
@@ -0,0 +1,54 @@
+import torch
+import math
+
+
+class MaternCovariance(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x1, x2, lengthscale, nu, dist_func):
+        if any(ctx.needs_input_grad[:2]):
+            raise RuntimeError("MaternCovariance cannot compute gradients with "
+                               "respect to x1 and x2")
+        if lengthscale.size(-1) > 1:
+            raise ValueError("MaternCovariance cannot handle multiple lengthscales")
+        # Subtract mean for numerical stability. Won't affect computations
+        # because covariance matrix is stationary.
+        needs_grad = any(ctx.needs_input_grad)
+        mean = x1.contiguous().view(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]
+        x1_ = (x1 - mean).div(lengthscale)
+        x2_ = (x2 - mean).div(lengthscale)
+        scaled_unitless_dist = dist_func(x1_, x2_).mul_(math.sqrt(2 * nu))
+        if nu == 0.5:
+            # 1 kernel sized Tensor if no grad else 2
+            scaled_unitless_dist_ = scaled_unitless_dist.clone() if needs_grad else scaled_unitless_dist
+            exp_component = scaled_unitless_dist_.neg_().exp_()
+            covar_mat = exp_component
+            if needs_grad:
+                d_output_d_input = scaled_unitless_dist.div_(lengthscale).mul_(exp_component)
+        elif nu == 1.5:
+            # 2 kernel sized Tensors if no grad else 3
+            if needs_grad:
+                scaled_unitless_dist_ = scaled_unitless_dist.clone()
+            linear_term = scaled_unitless_dist.clone().add_(1)
+            exp_component = scaled_unitless_dist.neg_().exp_()
+            covar_mat = linear_term.mul_(exp_component)
+            if needs_grad:
+                d_output_d_input = scaled_unitless_dist_.pow_(2).div_(lengthscale).mul_(exp_component)
+        elif nu == 2.5:
+            # 3 kernel sized Tensors if no grad else 4
+            linear_term = scaled_unitless_dist.clone().add_(1)
+            quadratic_term = scaled_unitless_dist.clone().pow_(2).div_(3)
+            exp_component = scaled_unitless_dist.neg_().exp_()
+            if needs_grad:
+                covar_mat = (linear_term + quadratic_term).mul_(exp_component)
+                d_output_d_input = linear_term.mul_(quadratic_term).mul_(exp_component).div_(lengthscale)
+            else:
+                covar_mat = exp_component.mul_(linear_term.add_(quadratic_term))
+        if needs_grad:
+            ctx.save_for_backward(d_output_d_input)
+        return covar_mat
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        d_output_d_input = ctx.saved_tensors[0]
+        lengthscale_grad = grad_output * d_output_d_input
+        return None, None, lengthscale_grad, None, None
diff --git a/gpytorch/functions/rbf_covariance.py b/gpytorch/functions/rbf_covariance.py
new file mode 100644
index 00000000..766678bf
--- /dev/null
+++ b/gpytorch/functions/rbf_covariance.py
@@ -0,0 +1,28 @@
+import torch
+
+
+class RBFCovariance(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x1, x2, lengthscale, sq_dist_func):
+        if any(ctx.needs_input_grad[:2]):
+            raise RuntimeError("RBFCovariance cannot compute gradients with "
+                               "respect to x1 and x2")
+        if lengthscale.size(-1) > 1:
+            raise ValueError("RBFCovariance cannot handle multiple lengthscales")
+        needs_grad = any(ctx.needs_input_grad)
+        x1_ = x1.div(lengthscale)
+        x2_ = x2.div(lengthscale)
+        unitless_sq_dist = sq_dist_func(x1_, x2_)
+        # clone because inplace operations will mess with what's saved for backward
+        unitless_sq_dist_ = unitless_sq_dist.clone() if needs_grad else unitless_sq_dist
+        covar_mat = unitless_sq_dist_.div_(-2.).exp_()
+        if needs_grad:
+            d_output_d_input = unitless_sq_dist.mul_(covar_mat).div_(lengthscale)
+            ctx.save_for_backward(d_output_d_input)
+        return covar_mat
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        d_output_d_input = ctx.saved_tensors[0]
+        lengthscale_grad = grad_output * d_output_d_input
+        return None, None, lengthscale_grad, None
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index 2e9d2390..3eb58557 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -21,7 +21,7 @@ class Distance(torch.jit.ScriptModule):
         self._postprocess = postprocess_script
 
     @torch.jit.script_method
-    def _jit_sq_dist(self, x1, x2, diag, x1_eq_x2):
+    def _jit_sq_dist(self, x1, x2, x1_eq_x2, postprocess):
         # Compute squared distance matrix using quadratic expansion
         x1_norm = x1.pow(2).sum(dim=-1, keepdim=True)
         if bool(x1_eq_x2):
@@ -38,14 +38,15 @@ class Distance(torch.jit.ScriptModule):
 
         # Zero out negative values
         res.clamp_min_(0)
-
-        return self._postprocess(res)
+        return self._postprocess(res) if bool(postprocess) else res
 
     @torch.jit.script_method
-    def _jit_dist(self, x1, x2, diag, x1_eq_x2):
-        res = self._jit_sq_dist(x1, x2, diag, x1_eq_x2)
+    def _jit_dist(self, x1, x2, x1_eq_x2, postprocess, false_tensor):
+        # Need to do a hack to get a False tensor because pytorch stable doesn't
+        # support creating of tensors in torch scripts
+        res = self._jit_sq_dist(x1, x2, x1_eq_x2, postprocess=false_tensor)
         res = res.clamp_min_(1e-30).sqrt_()
-        return self._postprocess(res)
+        return self._postprocess(res) if bool(postprocess) else res
 
 
 class Kernel(Module):
@@ -218,10 +219,6 @@ class Kernel(Module):
         """
         raise NotImplementedError()
 
-    @torch.jit.script_method
-    def _postprocess(self, dist_mat):
-        return dist_mat
-
     def __getstate__(self):
         # JIT ScriptModules cannot be pickled
         self.distance_module = None
@@ -238,6 +235,7 @@ class Kernel(Module):
         batch_dims=None,
         square_dist=False,
         dist_postprocess_func=default_postprocess_script,
+        postprocess=True,
         **params
     ):
         """
@@ -269,11 +267,13 @@ class Kernel(Module):
             x2 = x2.unsqueeze(0).transpose(0, -1)
 
         x1_eq_x2 = torch.equal(x1, x2)
+        # torch scripts expect tensors
+        postprocess = torch.tensor(postprocess)
 
         res = None
 
         # Cache the Distance object or else JIT will recompile every time
-        if not self.distance_module:
+        if not self.distance_module or self.distance_module._postprocess != dist_postprocess_func:
             self.distance_module = Distance(dist_postprocess_func)
 
         if diag:
@@ -284,18 +284,20 @@ class Kernel(Module):
                     res = torch.zeros(x1.shape[0] * x1.shape[-1], x2.shape[-2], dtype=x1.dtype, device=x1.device)
                 else:
                     res = torch.zeros(x1.shape[0], x2.shape[-2], dtype=x1.dtype, device=x1.device)
-                return dist_postprocess_func(res)
+                if postprocess:
+                    res = dist_postprocess_func(res)
+                return res
             else:
+                res = torch.norm(x1 - x2, p=2, dim=-1)
                 if square_dist:
-                    res = (x1 - x2).pow(2).sum(-1)
-                else:
-                    res = torch.norm(x1 - x2, p=2, dim=-1)
+                    res = res.pow(2)
+            if postprocess:
+                res = dist_postprocess_func(res)
 
-            res = dist_postprocess_func(res)
-        elif not square_dist:
-            res = self.distance_module._jit_dist(x1, x2, torch.tensor(diag), torch.tensor(x1_eq_x2))
+        elif square_dist:
+            res = self.distance_module._jit_sq_dist(x1, x2, torch.tensor(x1_eq_x2), postprocess)
         else:
-            res = self.distance_module._jit_sq_dist(x1, x2, torch.tensor(diag), torch.tensor(x1_eq_x2))
+            res = self.distance_module._jit_dist(x1, x2, torch.tensor(x1_eq_x2), postprocess, torch.tensor(False))
 
         if batch_dims == (0, 2):
             if diag:
diff --git a/gpytorch/kernels/matern_kernel.py b/gpytorch/kernels/matern_kernel.py
index 5925c4ae..b838ea61 100644
--- a/gpytorch/kernels/matern_kernel.py
+++ b/gpytorch/kernels/matern_kernel.py
@@ -4,6 +4,7 @@ import math
 import torch
 from .kernel import Kernel
 from torch.nn.functional import softplus
+from ..functions import MaternCovariance
 
 
 class MaternKernel(Kernel):
@@ -106,17 +107,25 @@ class MaternKernel(Kernel):
         self.nu = nu
 
     def forward(self, x1, x2, **params):
-        mean = x1.contiguous().view(-1, 1, x1.size(-1)).mean(0, keepdim=True)
-
-        x1_ = (x1 - mean).div(self.lengthscale)
-        x2_ = (x2 - mean).div(self.lengthscale)
-        distance = self._covar_dist(x1_, x2_, **params)
-        exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance)
-
-        if self.nu == 0.5:
-            constant_component = 1
-        elif self.nu == 1.5:
-            constant_component = (math.sqrt(3) * distance).add(1)
-        elif self.nu == 2.5:
-            constant_component = (math.sqrt(5) * distance).add(1).add(5.0 / 3.0 * distance ** 2)
-        return constant_component * exp_component
+        if (
+            x1.requires_grad
+            or x2.requires_grad
+            or (self.ard_num_dims is not None and self.ard_num_dims > 1)
+            or params.get('diag', False)
+        ):
+            mean = x1.contiguous().view(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]
+
+            x1_ = (x1 - mean).div(self.lengthscale)
+            x2_ = (x2 - mean).div(self.lengthscale)
+            distance = self._covar_dist(x1_, x2_, **params)
+            exp_component = torch.exp(-math.sqrt(self.nu * 2) * distance)
+
+            if self.nu == 0.5:
+                constant_component = 1
+            elif self.nu == 1.5:
+                constant_component = (math.sqrt(3) * distance).add(1)
+            elif self.nu == 2.5:
+                constant_component = (math.sqrt(5) * distance).add(1).add(5.0 / 3.0 * distance ** 2)
+            return constant_component * exp_component
+        return MaternCovariance().apply(x1, x2, self.lengthscale, self.nu,
+                                        lambda x1, x2: self._covar_dist(x1, x2, **params))
diff --git a/gpytorch/kernels/rbf_kernel.py b/gpytorch/kernels/rbf_kernel.py
index 522cbe8d..ce5dc083 100644
--- a/gpytorch/kernels/rbf_kernel.py
+++ b/gpytorch/kernels/rbf_kernel.py
@@ -3,6 +3,7 @@
 from .kernel import Kernel
 from torch.nn.functional import softplus
 import torch
+from ..functions import RBFCovariance
 
 
 @torch.jit.script
@@ -94,7 +95,21 @@ class RBFKernel(Kernel):
         )
 
     def forward(self, x1, x2, diag=False, **params):
-        x1_ = x1.div(self.lengthscale)
-        x2_ = x2.div(self.lengthscale)
-        diff = self._covar_dist(x1_, x2_, square_dist=True, diag=diag, dist_postprocess_func=postprocess_rbf, **params)
-        return diff
+        if (
+            x1.requires_grad
+            or x2.requires_grad
+            or (self.ard_num_dims is not None and self.ard_num_dims > 1)
+            or diag
+        ):
+            x1_ = x1.div(self.lengthscale)
+            x2_ = x2.div(self.lengthscale)
+            return self._covar_dist(x1_, x2_, square_dist=True, diag=diag,
+                                    dist_postprocess_func=postprocess_rbf,
+                                    postprocess=True, **params)
+        return RBFCovariance().apply(x1, x2, self.lengthscale,
+                                     lambda x1, x2: self._covar_dist(x1, x2,
+                                                                     square_dist=True,
+                                                                     diag=False,
+                                                                     dist_postprocess_func=postprocess_rbf,
+                                                                     postprocess=False,
+                                                                     **params))


Test patch:
diff --git a/test/functions/test_matern_covariance.py b/test/functions/test_matern_covariance.py
new file mode 100644
index 00000000..fc6b54dc
--- /dev/null
+++ b/test/functions/test_matern_covariance.py
@@ -0,0 +1,100 @@
+#!/usr/bin/env python3
+
+import math
+import unittest
+import torch
+import gpytorch
+
+
+def dist_func(x1, x2):
+    dist_module = gpytorch.kernels.kernel.Distance()
+    return dist_module._jit_dist(x1, x2, torch.tensor(torch.equal(x1, x2)),
+                                 postprocess=torch.tensor(False),
+                                 false_tensor=torch.tensor(False))
+
+
+class TestMaternCovariance(unittest.TestCase):
+    def test_1_2_forward(self):
+        nu = 1 / 2
+        batch_size = (3, 2, 4)
+        x1 = torch.randn(*batch_size, 7, 9)
+        x2 = torch.randn(*batch_size, 6, 9)
+        # Doesn't support ARD
+        lengthscale = torch.randn(*batch_size).view(*batch_size, 1, 1) ** 2
+
+        res = gpytorch.functions.MaternCovariance().apply(x1, x2, lengthscale, nu, dist_func)
+        scaled_unitless_dist = math.sqrt(nu * 2) * dist_func(x1, x2).div(lengthscale)
+        exp_component = torch.exp(-scaled_unitless_dist)
+        actual = exp_component
+        self.assertTrue(torch.allclose(res, actual))
+
+    def test_1_2_backward(self):
+        nu = 1 / 2
+        batch_size = (3, 2, 4)
+        x1 = torch.randn(*batch_size, 7, 9, dtype=torch.float64)
+        x2 = torch.randn(*batch_size, 6, 9, dtype=torch.float64)
+        lengthscale = torch.randn(
+            *batch_size, dtype=torch.float64, requires_grad=True).view(*batch_size, 1, 1) ** 2
+        f = lambda x1, x2, l: gpytorch.functions.MaternCovariance().apply(x1, x2, l, nu, dist_func)
+        try:
+            torch.autograd.gradcheck(f, (x1, x2, lengthscale))
+        except RuntimeError:
+            self.fail(f"Gradcheck failed")
+
+    def test_3_2_forward(self):
+        nu = 3 / 2
+        batch_size = (3, 2, 4)
+        x1 = torch.randn(*batch_size, 7, 9)
+        x2 = torch.randn(*batch_size, 6, 9)
+        # Doesn't support ARD
+        lengthscale = torch.randn(*batch_size).view(*batch_size, 1, 1) ** 2
+
+        res = gpytorch.functions.MaternCovariance().apply(x1, x2, lengthscale, nu, dist_func)
+        scaled_unitless_dist = math.sqrt(nu * 2) * dist_func(x1, x2).div(lengthscale)
+        exp_component = torch.exp(-scaled_unitless_dist)
+        actual = exp_component * (1 + scaled_unitless_dist)
+        self.assertTrue(torch.allclose(res, actual))
+
+    def test_3_2_backward(self):
+        nu = 3 / 2
+        batch_size = (3, 2, 4)
+        x1 = torch.randn(*batch_size, 7, 9, dtype=torch.float64)
+        x2 = torch.randn(*batch_size, 6, 9, dtype=torch.float64)
+        lengthscale = torch.randn(
+            *batch_size, dtype=torch.float64, requires_grad=True).view(*batch_size, 1, 1) ** 2
+        f = lambda x1, x2, l: gpytorch.functions.MaternCovariance().apply(x1, x2, l, nu, dist_func)
+        try:
+            torch.autograd.gradcheck(f, (x1, x2, lengthscale))
+        except RuntimeError:
+            self.fail(f"Gradcheck failed")
+
+    def test_5_2_forward(self):
+        nu = 5 / 2
+        batch_size = (3, 2, 4)
+        x1 = torch.randn(*batch_size, 7, 9)
+        x2 = torch.randn(*batch_size, 6, 9)
+        # Doesn't support ARD
+        lengthscale = torch.randn(*batch_size).view(*batch_size, 1, 1) ** 2
+
+        res = gpytorch.functions.MaternCovariance().apply(x1, x2, lengthscale, nu, dist_func)
+        scaled_unitless_dist = math.sqrt(nu * 2) * dist_func(x1, x2).div(lengthscale)
+        exp_component = torch.exp(-scaled_unitless_dist)
+        actual = exp_component * (1 + scaled_unitless_dist + scaled_unitless_dist ** 2 / 3)
+        self.assertTrue(torch.allclose(res, actual))
+
+    def test_5_2_backward(self):
+        nu = 5 / 2
+        batch_size = (3, 2, 4)
+        x1 = torch.randn(*batch_size, 7, 9, dtype=torch.float64)
+        x2 = torch.randn(*batch_size, 6, 9, dtype=torch.float64)
+        lengthscale = torch.randn(
+            *batch_size, dtype=torch.float64, requires_grad=True).view(*batch_size, 1, 1) ** 2
+        f = lambda x1, x2, l: gpytorch.functions.MaternCovariance().apply(x1, x2, l, nu, dist_func)
+        try:
+            torch.autograd.gradcheck(f, (x1, x2, lengthscale))
+        except RuntimeError:
+            self.fail(f"Gradcheck failed")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/functions/test_rbf_covariance.py b/test/functions/test_rbf_covariance.py
new file mode 100644
index 00000000..127bc157
--- /dev/null
+++ b/test/functions/test_rbf_covariance.py
@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+
+import torch
+import unittest
+import gpytorch
+
+
+def sq_dist_func(x1, x2):
+    dist_module = gpytorch.kernels.kernel.Distance()
+    return dist_module._jit_sq_dist(x1, x2, torch.tensor(torch.equal(x1, x2)),
+                                    postprocess=torch.tensor(False))
+
+
+class TestRBFCovariance(unittest.TestCase):
+    def test_forward(self):
+        batch_size = (3, 2, 4)
+        x1 = torch.randn(*batch_size, 7, 9)
+        x2 = torch.randn(*batch_size, 6, 9)
+        # Doesn't support ARD
+        lengthscale = torch.randn(*batch_size).view(*batch_size, 1, 1) ** 2
+        res = gpytorch.functions.RBFCovariance().apply(x1, x2, lengthscale, sq_dist_func)
+        actual = sq_dist_func(x1, x2).div(-2 * lengthscale**2).exp()
+        self.assertTrue(torch.allclose(res, actual))
+
+    def test_backward(self):
+        batch_size = (3, 2, 4)
+        x1 = torch.randn(*batch_size, 7, 9, dtype=torch.float64)
+        x2 = torch.randn(*batch_size, 6, 9, dtype=torch.float64)
+        lengthscale = torch.randn(*batch_size, dtype=torch.float64, requires_grad=True).view(*batch_size, 1, 1) ** 2
+        f = lambda x1, x2, l: gpytorch.functions.RBFCovariance().apply(x1, x2, l, sq_dist_func)
+        try:
+            torch.autograd.gradcheck(f, (x1, x2, lengthscale))
+        except RuntimeError:
+            self.fail("Gradcheck failed")
+
+
+if __name__ == "__main__":
+    unittest.main()