Repository: cornellius-gp/gpytorch
Issue: Grid kernel has trouble with 1D inputs

I'm new to both PyTorch and GPyTorch, so please forgive me if I'm doing something silly.


The GridKernel seems to have problems with a 1-dimensional input, where 1-dimensional and 2-dimensional inputs are handled differently for some reason.  2- (and 3-) dimensional cases run without the error below.

Working example:
From the current master branch: 
Modify the test_grid_gp_regression.py script to create a similar script, here called test_1D_grid_gp_regression.py, where the diff is this:
```
$ diff test_grid_gp_regression.py test_1D_grid_gp_regression.py
15c15
<     train_y = torch.sin((train_x[:, 0] + train_x[:, 1]) * (2 * math.pi)) + torch.randn_like(train_x[:, 0]).mul(0.01)
---
>     train_y = torch.sin((train_x[:, 0]) * (2 * math.pi)) + torch.randn_like(train_x[:, 0]).mul(0.01)
58c58
<         grid_bounds = [(0, 1), (0, 2)]
---
>         grid_bounds = [(0, 1)]
```

I get the following error when running the test:
```
$ python test_1D_grid_gp_regression.py 
E.
======================================================================
ERROR: test_grid_gp_mean_abs_error (__main__.TestGridGPRegression)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "test_1D_grid_gp_regression.py", line 87, in test_grid_gp_mean_abs_error
    loss.backward()
  File "/usr/local/lib/python3.6/dist-packages/torch/tensor.py", line 102, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/usr/local/lib/python3.6/dist-packages/torch/autograd/__init__.py", line 90, in backward
    allow_unreachable=True)  # allow_unreachable flag
  File "/home/gpytorch/gpytorch/functions/_matmul.py", line 34, in backward
    rhs = self.saved_tensors[0]
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation

----------------------------------------------------------------------
Ran 2 tests in 0.088s

FAILED (errors=1)
```


Changed files:
test/examples/test_grid_gp_regression.py, gpytorch/functions/_inv_quad_log_det.py, gpytorch/lazy/lazy_tensor.py, gpytorch/functions/_inv_matmul.py

diff --git a/gpytorch/functions/_inv_matmul.py b/gpytorch/functions/_inv_matmul.py
index 527729c5..959f9a9d 100644
--- a/gpytorch/functions/_inv_matmul.py
+++ b/gpytorch/functions/_inv_matmul.py
@@ -6,9 +6,8 @@ from .. import settings
 
 
 class InvMatmul(Function):
-    def __init__(self, representation_tree, preconditioner=None, has_left=False):
+    def __init__(self, representation_tree, has_left=False):
         self.representation_tree = representation_tree
-        self.preconditioner = preconditioner
         self.has_left = has_left
 
     def forward(self, *args):
@@ -22,6 +21,9 @@ class InvMatmul(Function):
         orig_right_tensor = right_tensor
         lazy_tsr = self.representation_tree(*matrix_args)
 
+        with torch.no_grad():
+            self.preconditioner = lazy_tsr.detach()._inv_matmul_preconditioner()
+
         self.is_vector = False
         if right_tensor.ndimension() == 1:
             right_tensor = right_tensor.unsqueeze(-1)
diff --git a/gpytorch/functions/_inv_quad_log_det.py b/gpytorch/functions/_inv_quad_log_det.py
index 14e5cf54..336bd69e 100644
--- a/gpytorch/functions/_inv_quad_log_det.py
+++ b/gpytorch/functions/_inv_quad_log_det.py
@@ -24,8 +24,6 @@ class InvQuadLogDet(Function):
         batch_shape=torch.Size(),
         inv_quad=False,
         logdet=False,
-        preconditioner=None,
-        logdet_correction=None,
         probe_vectors=None,
         probe_vector_norms=None,
     ):
@@ -39,8 +37,6 @@ class InvQuadLogDet(Function):
         self.batch_shape = batch_shape
         self.inv_quad = inv_quad
         self.logdet = logdet
-        self.preconditioner = preconditioner
-        self.logdet_correction = logdet_correction
 
         if (probe_vectors is None or probe_vector_norms is None) and logdet:
             num_random_probes = settings.num_trace_samples.value()
@@ -75,6 +71,8 @@ class InvQuadLogDet(Function):
 
         # Get closure for matmul
         lazy_tsr = self.representation_tree(*matrix_args)
+        with torch.no_grad():
+            preconditioner, logdet_correction = lazy_tsr.detach()._preconditioner()
 
         # Collect terms for LinearCG
         # We use LinearCG for both matrix solves and for stochastically estimating the log det
@@ -100,10 +98,10 @@ class InvQuadLogDet(Function):
         rhs = torch.cat(rhs_list, -1)
         t_mat = None
         if self.logdet and settings.skip_logdet_forward.off():
-            solves, t_mat = lazy_tsr._solve(rhs, self.preconditioner, num_tridiag=num_random_probes)
+            solves, t_mat = lazy_tsr._solve(rhs, preconditioner, num_tridiag=num_random_probes)
 
         else:
-            solves = lazy_tsr._solve(rhs, self.preconditioner, num_tridiag=0)
+            solves = lazy_tsr._solve(rhs, preconditioner, num_tridiag=0)
 
         # Final values to return
         logdet_term = torch.zeros(lazy_tsr.batch_shape, dtype=self.dtype, device=self.device)
@@ -121,8 +119,8 @@ class InvQuadLogDet(Function):
                 logdet_term, = slq.evaluate(self.matrix_shape, eigenvalues, eigenvectors, [lambda x: x.log()])
 
                 # Add correction
-                if self.logdet_correction is not None:
-                    logdet_term = logdet_term + self.logdet_correction
+                if logdet_correction is not None:
+                    logdet_term = logdet_term + logdet_correction
 
         # Extract inv_quad solves from all the solves
         if self.inv_quad:
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
index 3d85bd13..d362a261 100644
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -681,7 +681,6 @@ class LazyTensor(object):
 
         func = InvMatmul(
             self.representation_tree(),
-            preconditioner=self._inv_matmul_preconditioner(),
             has_left=(left_tensor is not None),
         )
         if left_tensor is None:
@@ -758,8 +757,6 @@ class LazyTensor(object):
             device=self.device,
             inv_quad=(inv_quad_rhs is not None),
             logdet=logdet,
-            preconditioner=self._preconditioner()[0],
-            logdet_correction=self._preconditioner()[1],
             probe_vectors=probe_vectors,
             probe_vector_norms=probe_vector_norms,
         )(*args)
diff --git a/test/examples/test_grid_gp_regression.py b/test/examples/test_grid_gp_regression.py
index 6e53e1b8..02787ab8 100644
--- a/test/examples/test_grid_gp_regression.py
+++ b/test/examples/test_grid_gp_regression.py
@@ -12,14 +12,14 @@ from torch import optim
 
 def make_data(grid, cuda=False):
     train_x = gpytorch.utils.grid.create_data_from_grid(grid)
-    train_y = torch.sin((train_x[:, 0] + train_x[:, 1]) * (2 * math.pi)) + torch.randn_like(train_x[:, 0]).mul(0.01)
+    train_y = torch.sin((train_x.sum(-1)) * (2 * math.pi)) + torch.randn_like(train_x[:, 0]).mul(0.01)
     n = 20
     test_x = torch.zeros(int(pow(n, 2)), 2)
     for i in range(n):
         for j in range(n):
             test_x[i * n + j][0] = float(i) / (n - 1)
             test_x[i * n + j][1] = float(j) / (n - 1)
-    test_y = torch.sin(((test_x[:, 0] + test_x[:, 1]) * (2 * math.pi)))
+    test_y = torch.sin(((test_x.sum(-1)) * (2 * math.pi)))
     if cuda:
         train_x = train_x.cuda()
         train_y = train_y.cuda()
@@ -53,9 +53,9 @@ class TestGridGPRegression(unittest.TestCase):
         if hasattr(self, "rng_state"):
             torch.set_rng_state(self.rng_state)
 
-    def test_grid_gp_mean_abs_error(self, cuda=False):
+    def test_grid_gp_mean_abs_error(self, num_dim=1, cuda=False):
         device = torch.device("cuda") if cuda else torch.device("cpu")
-        grid_bounds = [(0, 1), (0, 2)]
+        grid_bounds = [(0, 1)] if num_dim == 1 else [(0, 1), (0, 2)]
         grid_size = 25
         grid = torch.zeros(grid_size, len(grid_bounds), device=device)
         for i in range(len(grid_bounds)):
@@ -108,6 +108,13 @@ class TestGridGPRegression(unittest.TestCase):
         if torch.cuda.is_available():
             self.test_grid_gp_mean_abs_error(cuda=True)
 
+    def test_grid_gp_mean_abs_error_2d(self):
+        self.test_grid_gp_mean_abs_error(num_dim=2)
+
+    def test_grid_gp_mean_abs_error_2d_cuda(self):
+        if torch.cuda.is_available():
+            self.test_grid_gp_mean_abs_error(cuda=True, num_dim=2)
+
 
 if __name__ == "__main__":
     unittest.main()