Repository: cornellius-gp/gpytorch
Issue: [Bug] psd_safe_cholesky less stable with small noise

# 🐛 Bug

I'm getting errors (see below) now from `psd_safe_cholesky` which weren't there a few commits back.

* Current commit has this error: 6a235fb6ec339e77a86a3cb20c715071c71ac9fc
* A recent commit without this error: 308a18094d31d5cd2fbb7d942212669141dd8cec. I'm not sure what the *most* recent commit without this error is, though.

## To reproduce

```python
import torch
import gpytorch

device = "cuda"

train_x = torch.tensor([
    [-55.5527, -42.9031],
    [ 63.1439, -36.1096],
    [-44.4039, -51.3083],
    [-37.9279,  44.9540],
    [-54.8894,   1.7219],
    [-26.8499, -40.2266],
    [ 17.5439,  26.0574],
    [-31.0171,  -2.4689],
    [ 37.8813, -39.7865],
    [-10.5718,  67.5785],
    [ 13.2326, -45.9986],
    [ 51.3409, -21.4855],
    [ 64.9224, -52.7161],
    [ 13.1864,  63.9087],
    [-67.8713, -59.4553],
    [  3.3758,  13.9108],
    [ 12.1702,  35.1921],
    [ 47.5776,  69.5157],
    [ -7.5170,  59.2905],
    [-34.0096, -51.0971],
    [-43.5709,  14.1764],
    [-17.3790,  57.2735],
    [ 46.1217,   4.4776],
    [ 17.8831,  68.3338],
    [ -0.1228,  61.4203],
    [ 46.7972,  27.2138],
    [-35.2764, -19.0872],
    [ 36.9366,  45.3593],
    [-20.7411,  40.9940],
    [ 16.2726, -29.3061]
], device=device)
train_y = torch.tensor([59.1276,  -53.7113,   -1.7330, -105.8712,   10.8977,  -16.1213,
         -43.1551,   55.5215,   32.3309,   87.8219,    4.2356,   28.3534,
         -60.4486,  115.4285,   68.5259,  -64.0631,  -12.9961,   34.0834,
          61.4407,  -28.7934,  -31.1267,   18.5751,  -70.5619,  121.7277,
          90.2278,   73.1477,   36.8967,   48.5388,  -68.2077,    1.3766], device=device)

class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y):
        likelihood = gpytorch.likelihoods.GaussianLikelihood()
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.MaternKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

def train(model, train_x, train_y, train_steps):
    model.set_train_data(train_x, train_y, strict=False)
    model.train()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

    mll = gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model)

    for i in range(train_steps):
        optimizer.zero_grad()
        output = model(train_x)
        loss = -mll(output, train_y)
        loss.backward()
        optimizer.step()
    model.eval()

model = ExactGPModel(train_x, train_y)
model.to(device)
model.likelihood.initialize(noise=1e-5)
model.likelihood.noise_covar.raw_noise.requires_grad_(False)
train(model, train_x, train_y, train_steps=100)
```

## Stack traces

If `device = "cpu"`"

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-5-dcd721edef0b> in <module>
     73 model.likelihood.initialize(noise=1e-5)
     74 model.likelihood.noise_covar.raw_noise.requires_grad_(False)
---> 75 train(model, train_x, train_y, train_steps=100)

<ipython-input-5-dcd721edef0b> in train(model, train_x, train_y, train_steps)
     64         optimizer.zero_grad()
     65         output = model(train_x)
---> 66         loss = -mll(output, train_y)
     67         loss.backward()
     68         optimizer.step()

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
     20 
     21     def __call__(self, *inputs, **kwargs):
---> 22         outputs = self.forward(*inputs, **kwargs)
     23         if isinstance(outputs, list):
     24             return [_validate_module_outputs(output) for output in outputs]

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py in forward(self, output, target, *params)
     26         # Get the log prob of the marginal distribution
     27         output = self.likelihood(output, *params)
---> 28         res = output.log_prob(target)
     29 
     30         # Add terms for SGPR / when inducing points are learned

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/distributions/multivariate_normal.py in log_prob(self, value)
    127 
    128         # Get log determininat and first part of quadratic form
--> 129         inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
    130 
    131         res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)
    990             from .chol_lazy_tensor import CholLazyTensor
    991 
--> 992             cholesky = CholLazyTensor(self.cholesky())
    993             return cholesky.inv_quad_logdet(inv_quad_rhs=inv_quad_rhs, logdet=logdet, reduce_inv_quad=reduce_inv_quad)
    994 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in cholesky(self, upper)
    716             (LazyTensor) Cholesky factor (lower triangular)
    717         """
--> 718         res = self._cholesky()
    719         if upper:
    720             res = res.transpose(-1, -2)

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/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 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in _cholesky(self)
    401             evaluated_mat.register_hook(_ensure_symmetric_grad)
    402 
--> 403         cholesky = psd_safe_cholesky(evaluated_mat.double()).to(self.dtype)
    404         return NonLazyTensor(cholesky)
    405 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/cholesky.py in psd_safe_cholesky(A, upper, out, jitter)
     44                 continue
     45 
---> 46         raise e
     47 
     48 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/cholesky.py in psd_safe_cholesky(A, upper, out, jitter)
     19     """
     20     try:
---> 21         L = torch.cholesky(A, upper=upper, out=out)
     22         # TODO: Remove once fixed in pytorch (#16780)
     23         if A.dim() > 2 and A.is_cuda:

RuntimeError: Lapack Error in potrf : the leading minor of order 1 is not positive definite at /pytorch/aten/src/TH/generic/THTensorLapack.cpp:658
```

If `device = "cuda"`

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-6-2b962e725013> in <module>
     73 model.likelihood.initialize(noise=1e-5)
     74 model.likelihood.noise_covar.raw_noise.requires_grad_(False)
---> 75 train(model, train_x, train_y, train_steps=100)

<ipython-input-6-2b962e725013> in train(model, train_x, train_y, train_steps)
     64         optimizer.zero_grad()
     65         output = model(train_x)
---> 66         loss = -mll(output, train_y)
     67         loss.backward()
     68         optimizer.step()

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/module.py in __call__(self, *inputs, **kwargs)
     20 
     21     def __call__(self, *inputs, **kwargs):
---> 22         outputs = self.forward(*inputs, **kwargs)
     23         if isinstance(outputs, list):
     24             return [_validate_module_outputs(output) for output in outputs]

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/mlls/exact_marginal_log_likelihood.py in forward(self, output, target, *params)
     26         # Get the log prob of the marginal distribution
     27         output = self.likelihood(output, *params)
---> 28         res = output.log_prob(target)
     29 
     30         # Add terms for SGPR / when inducing points are learned

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/distributions/multivariate_normal.py in log_prob(self, value)
    127 
    128         # Get log determininat and first part of quadratic form
--> 129         inv_quad, logdet = covar.inv_quad_logdet(inv_quad_rhs=diff.unsqueeze(-1), logdet=True)
    130 
    131         res = -0.5 * sum([inv_quad, logdet, diff.size(-1) * math.log(2 * math.pi)])

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in inv_quad_logdet(self, inv_quad_rhs, logdet, reduce_inv_quad)
    990             from .chol_lazy_tensor import CholLazyTensor
    991 
--> 992             cholesky = CholLazyTensor(self.cholesky())
    993             return cholesky.inv_quad_logdet(inv_quad_rhs=inv_quad_rhs, logdet=logdet, reduce_inv_quad=reduce_inv_quad)
    994 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in cholesky(self, upper)
    716             (LazyTensor) Cholesky factor (lower triangular)
    717         """
--> 718         res = self._cholesky()
    719         if upper:
    720             res = res.transpose(-1, -2)

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/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 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/lazy/lazy_tensor.py in _cholesky(self)
    401             evaluated_mat.register_hook(_ensure_symmetric_grad)
    402 
--> 403         cholesky = psd_safe_cholesky(evaluated_mat.double()).to(self.dtype)
    404         return NonLazyTensor(cholesky)
    405 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/cholesky.py in psd_safe_cholesky(A, upper, out, jitter)
     44                 continue
     45 
---> 46         raise e
     47 
     48 

~/.cache/pypoetry/virtualenvs/bayes-optim-py3.7/src/gpytorch/gpytorch/utils/cholesky.py in psd_safe_cholesky(A, upper, out, jitter)
     19     """
     20     try:
---> 21         L = torch.cholesky(A, upper=upper, out=out)
     22         # TODO: Remove once fixed in pytorch (#16780)
     23         if A.dim() > 2 and A.is_cuda:

RuntimeError: MAGMA potrf : A(1,1) is 0, A cannot be factorized at /pytorch/aten/src/THC/generic/THCTensorMathMagma.cu:597
```

## Expected Behavior

Training is no less stable than it was at, e.g., commit 308a18094d31d5cd2fbb7d942212669141dd8cec.

## System information

**Please complete the following information:**
- GPyTorch Version: commit 6a235fb6ec339e77a86a3cb20c715071c71ac9fc
- PyTorch Version: 1.0.1.post2
- OS: Ubuntu 16.04.5

## Additional context

Removing the `model.likelihood` lines makes this work, so it only happens with small noise.

Changed files:
gpytorch/module.py, gpytorch/constraints/constraints.py, test/kernels/_base_kernel_test_case.py

diff --git a/gpytorch/constraints/constraints.py b/gpytorch/constraints/constraints.py
index 4b6ffc18..2aa8b3a4 100644
--- a/gpytorch/constraints/constraints.py
+++ b/gpytorch/constraints/constraints.py
@@ -44,7 +44,13 @@ class Interval(Module):
         return self._transform is not None
 
     def check(self, tensor):
-        return torch.all(tensor <= self.upper_bound) and torch.all(tensor >= self.lower_bound)
+        return bool(torch.all(tensor <= self.upper_bound) and torch.all(tensor >= self.lower_bound))
+
+    def check_raw(self, tensor):
+        return bool(
+            torch.all((self.transform(tensor) <= self.upper_bound))
+            and torch.all(self.transform(tensor) >= self.lower_bound)
+        )
 
     def intersect(self, other):
         """
diff --git a/gpytorch/module.py b/gpytorch/module.py
index 81ef8cee..3f01b8d7 100644
--- a/gpytorch/module.py
+++ b/gpytorch/module.py
@@ -77,12 +77,28 @@ class Module(nn.Module):
             elif name not in self._parameters:
                 setattr(self, name, val)
             elif torch.is_tensor(val):
+                constraint = self.constraint_for_parameter_name(name)
+                if constraint is not None and not constraint.check_raw(val):
+                    raise RuntimeError(
+                        "Attempting to manually set a parameter value that is out of bounds of "
+                        f"its current constraints, {constraint}. "
+                        "Most likely, you want to do the following:\\n likelihood = GaussianLikelihood"
+                        "(noise_constraint=gpytorch.constraints.GreaterThan(better_lower_bound))"
+                    )
                 try:
                     self.__getattr__(name).data.copy_(val.expand_as(self.__getattr__(name)))
                 except RuntimeError:
                     self.__getattr__(name).data.copy_(val.view_as(self.__getattr__(name)))
 
             elif isinstance(val, float):
+                constraint = self.constraint_for_parameter_name(name)
+                if constraint is not None and not constraint.check_raw(val):
+                    raise RuntimeError(
+                        "Attempting to manually set a parameter value that is out of bounds of "
+                        f"its current constraints, {constraint}. "
+                        "Most likely, you want to do the following:\\n likelihood = GaussianLikelihood"
+                        "(noise_constraint=gpytorch.constraints.GreaterThan(better_lower_bound))"
+                    )
                 self.__getattr__(name).data.fill_(val)
             else:
                 raise AttributeError("Type {t} not valid for initializing parameter {p}".format(t=type(val), p=name))
diff --git a/test/kernels/_base_kernel_test_case.py b/test/kernels/_base_kernel_test_case.py
index a7fbe9d4..0df9e725 100644
--- a/test/kernels/_base_kernel_test_case.py
+++ b/test/kernels/_base_kernel_test_case.py
@@ -20,7 +20,7 @@ class BaseKernelTestCase(object):
         kernel_basic = self.create_kernel_no_ard()
         covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().evaluate()
 
-        self.assertLess(torch.norm(covar_mat - covar_mat_actual), 1e-4)
+        self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)
 
     def test_active_dims_range(self):
         active_dims = list(range(3, 9))
@@ -30,7 +30,7 @@ class BaseKernelTestCase(object):
         kernel_basic = self.create_kernel_no_ard()
         covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().evaluate()
 
-        self.assertLess(torch.norm(covar_mat - covar_mat_actual), 1e-4)
+        self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)
 
     def test_no_batch_kernel_single_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard()
@@ -41,7 +41,7 @@ class BaseKernelTestCase(object):
         actual_mat_2 = kernel(x[1]).evaluate_kernel().evaluate()
         actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
 
     def test_single_batch_kernel_single_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
@@ -52,7 +52,7 @@ class BaseKernelTestCase(object):
         actual_mat_2 = kernel(x[1]).evaluate_kernel().evaluate()
         actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
 
     def test_no_batch_kernel_double_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([]))
@@ -68,7 +68,7 @@ class BaseKernelTestCase(object):
 
         actual_covar_mat = torch.cat([ac.unsqueeze(0) for ac in ij_actual_covars])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
 
     def test_no_batch_kernel_double_batch_x_ard(self):
         try:
@@ -88,7 +88,7 @@ class BaseKernelTestCase(object):
 
         actual_covar_mat = torch.cat([ac.unsqueeze(0) for ac in ij_actual_covars])
 
-        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat), 1e-4)
+        self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
 
     def test_smoke_double_batch_kernel_double_batch_x_no_ard(self):
         kernel = self.create_kernel_no_ard(batch_shape=torch.Size([3, 2]))