Repository: cornellius-gp/gpytorch
Pull request #510 (2019-02-07 20:20:18)

Stop cholesky_solve warnings on pytorch-nightly
This works on both pytorch stable and nightly. It creates a utility function that chooses the correct function to call (`torch.cholesky_solve` if avaialable, otherwise `torch.potrs`).

Also removes some stale tridiagonal cholesky code that is no longer used (and is very slow).


Patch:
diff --git a/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb b/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
index 5d67b429..ff84beaa 100644
--- a/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
+++ b/examples/01_Simple_GP_Regression/Simple_MultiGPU_GP_Regression.ipynb
@@ -73,14 +73,6 @@
     "scrolled": false
    },
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/jrg365/gpytorch/gpytorch/utils/pivoted_cholesky.py:103: UserWarning: torch.potrs is deprecated in favour of torch.cholesky_solve and will be removed in the next release. Please use torch.cholesky instead and note that the :attr:`upper` argument in torch.cholesky_solve defaults to ``False``.\n",
-      "  R = torch.potrs(low_rank_mat, torch.cholesky(shifted_mat, upper=True))\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
@@ -174,14 +166,6 @@
    "execution_count": 5,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/jrg365/gpytorch/gpytorch/utils/pivoted_cholesky.py:103: UserWarning: torch.potrs is deprecated in favour of torch.cholesky_solve and will be removed in the next release. Please use torch.cholesky instead and note that the :attr:`upper` argument in torch.cholesky_solve defaults to ``False``.\n",
-      "  R = torch.potrs(low_rank_mat, torch.cholesky(shifted_mat, upper=True))\n"
-     ]
-    }
    ],
    "source": [
     "# Get into evaluation (predictive posterior) mode\n",
diff --git a/examples/10_GP_Regression_Derivative_Information/Simple_GP_Regression_Derivative_Information_1d.ipynb b/examples/10_GP_Regression_Derivative_Information/Simple_GP_Regression_Derivative_Information_1d.ipynb
index 1c0e58b6..20d5d2c9 100644
--- a/examples/10_GP_Regression_Derivative_Information/Simple_GP_Regression_Derivative_Information_1d.ipynb
+++ b/examples/10_GP_Regression_Derivative_Information/Simple_GP_Regression_Derivative_Information_1d.ipynb
@@ -117,14 +117,6 @@
    "execution_count": 4,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/gpleiss/workspace/gpytorch/gpytorch/utils/pivoted_cholesky.py:101: UserWarning: torch.potrs is deprecated in favour of torch.cholesky_solve and will be removed in the next release. Please use torch.cholesky instead and note that the :attr:`upper` argument in torch.cholesky_solve defaults to ``False``.\n",
-      "  R = torch.potrs(low_rank_mat, torch.cholesky(shifted_mat, upper=True))\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
diff --git a/examples/10_GP_Regression_Derivative_Information/Simple_GP_Regression_Derivative_Information_2d.ipynb b/examples/10_GP_Regression_Derivative_Information/Simple_GP_Regression_Derivative_Information_2d.ipynb
index 2a043055..ac908310 100644
--- a/examples/10_GP_Regression_Derivative_Information/Simple_GP_Regression_Derivative_Information_2d.ipynb
+++ b/examples/10_GP_Regression_Derivative_Information/Simple_GP_Regression_Derivative_Information_2d.ipynb
@@ -141,14 +141,6 @@
    "execution_count": 5,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/gpleiss/workspace/gpytorch/gpytorch/utils/pivoted_cholesky.py:101: UserWarning: torch.potrs is deprecated in favour of torch.cholesky_solve and will be removed in the next release. Please use torch.cholesky instead and note that the :attr:`upper` argument in torch.cholesky_solve defaults to ``False``.\n",
-      "  R = torch.potrs(low_rank_mat, torch.cholesky(shifted_mat, upper=True))\n"
-     ]
-    },
     {
      "name": "stdout",
      "output_type": "stream",
diff --git a/gpytorch/models/exact_prediction_strategies.py b/gpytorch/models/exact_prediction_strategies.py
index 2e25eb21..914c398f 100644
--- a/gpytorch/models/exact_prediction_strategies.py
+++ b/gpytorch/models/exact_prediction_strategies.py
@@ -7,6 +7,7 @@ from ..distributions import MultivariateNormal
 from ..lazy import InterpolatedLazyTensor, LazyTensor, MatmulLazyTensor, RootLazyTensor, SumLazyTensor
 from ..utils.interpolation import left_interp, left_t_interp
 from ..utils.memoize import cached
+from ..utils.cholesky import cholesky_solve
 
 
 _PREDICTION_STRATEGY_REGISTRY = {}
@@ -141,10 +142,10 @@ class DefaultPredictionStrategy(object):
         small_system_rhs = targets - fant_mean - fant_train_covar.matmul(self.mean_cache)
         # Schur complement of a spd matrix is guaranteed to be positive definite
         if small_system_rhs.requires_grad or schur_complement.requires_grad:
-            # TODO: Delete this part of the if statement when PyTorch implements potrs derivative.
+            # TODO: Delete this part of the if statement when PyTorch implements cholesky_solve derivative.
             fant_cache_lower = torch.gesv(small_system_rhs.unsqueeze(-1), schur_complement)[0]
         else:
-            fant_cache_lower = torch.potrs(small_system_rhs, torch.cholesky(schur_complement, upper=True))
+            fant_cache_lower = cholesky_solve(small_system_rhs, torch.cholesky(schur_complement))
 
         # Get "a", the new upper portion of the cache corresponding to the old training points.
         fant_cache_upper = self.mean_cache.unsqueeze(-1) - fant_solve.matmul(fant_cache_lower)
@@ -196,10 +197,10 @@ class DefaultPredictionStrategy(object):
         # one of torch.svd, torch.qr, or torch.pinverse works in batch mode.
         cap_mat = new_root.transpose(-2, -1).matmul(new_root)
         if cap_mat.requires_grad or new_root.requires_grad:
-            # TODO: Delete this part of the if statement when PyTorch implements potrs derivative.
+            # TODO: Delete this part of the if statement when PyTorch implements cholesky_solve derivative.
             new_covar_cache = torch.gesv(new_root.transpose(-2, -1), cap_mat)[0].transpose(-2, -1)
         else:
-            new_covar_cache = torch.potrs(new_root.transpose(-2, -1), torch.cholesky(cap_mat, upper=True))
+            new_covar_cache = cholesky_solve(new_root.transpose(-2, -1), torch.cholesky(cap_mat))
             new_covar_cache = new_covar_cache.transpose(-2, -1)
 
         # Create new DefaultPredictionStrategy object
diff --git a/gpytorch/utils/cholesky.py b/gpytorch/utils/cholesky.py
index 66d84e44..e8d9ce17 100644
--- a/gpytorch/utils/cholesky.py
+++ b/gpytorch/utils/cholesky.py
@@ -1,90 +1,9 @@
 #!/usr/bin/env python3
 
 import warnings
-
 import torch
 
 
-def tridiag_batch_potrf(trid, upper=False):
-    """
-    """
-    if not torch.is_tensor(trid):
-        raise RuntimeError("tridiag_batch_potrf is only defined for tensors")
-
-    batch_size, diag_size, _ = trid.size()
-    batch_index = torch.arange(0, batch_size, dtype=torch.long, device=trid.device)
-    off_batch_index = batch_index.unsqueeze(1).repeat(diag_size - 1, 1).view(-1)
-    batch_index = batch_index.unsqueeze(1).repeat(diag_size, 1).view(-1)
-    diag_index = torch.arange(0, diag_size, dtype=torch.long, device=trid.device)
-    diag_index = diag_index.unsqueeze(1).repeat(1, batch_size).view(-1)
-    off_diag_index = torch.arange(0, diag_size - 1, dtype=torch.long, device=trid.device)
-    off_diag_index = off_diag_index.unsqueeze(1).repeat(1, batch_size).view(-1)
-
-    t_main_diag = trid[batch_index, diag_index, diag_index].view(diag_size, batch_size)
-    t_off_diag = trid[off_batch_index, off_diag_index + 1, off_diag_index].view(diag_size - 1, batch_size)
-
-    chol_main_diag = torch.empty_like(t_main_diag)
-    chol_off_diag = torch.empty_like(t_off_diag)
-
-    chol_main_diag[0].copy_(t_main_diag[0].sqrt())
-    for i in range(1, diag_size):
-        chol_off_diag[i - 1].copy_(t_off_diag[i - 1] / chol_main_diag[i - 1])
-        sq_value = t_main_diag[i] - chol_off_diag[i - 1] ** 2
-        chol_main_diag[i].copy_(torch.sqrt(sq_value))
-
-    res = torch.zeros_like(trid)
-    main_flattened_indices = batch_index * (batch_size * diag_size) + diag_index * (diag_size + 1)
-    off_flattened_indices = sum(
-        [off_batch_index * (batch_size * (diag_size - 1)), (off_diag_index + 1) * diag_size, off_diag_index]
-    )
-    res.view(-1).index_copy_(0, main_flattened_indices, chol_main_diag.view(-1))
-    res.view(-1).index_copy_(0, off_flattened_indices, chol_off_diag.view(-1))
-
-    if upper:
-        res = res.transpose(-1, -2)
-    return res
-
-
-def tridiag_batch_potrs(tensor, chol_trid, upper=True):
-    """
-    """
-    if not torch.is_tensor(chol_trid):
-        raise RuntimeError("tridiag_batch_potrf is only defined for tensors")
-
-    if not tensor.ndimension() == 3:
-        raise RuntimeError("Tensor should be 3 dimensional")
-
-    batch_size, diag_size, _ = chol_trid.size()
-    batch_index = torch.arange(0, batch_size, dtype=torch.long, device=tensor.device)
-    off_batch_index = batch_index.unsqueeze(1).repeat(1, diag_size - 1).view(-1)
-    batch_index = batch_index.unsqueeze(1).repeat(1, diag_size).view(-1)
-    diag_index = torch.arange(0, diag_size, dtype=torch.long, device=tensor.device)
-    diag_index = diag_index.unsqueeze(1).repeat(batch_size, 1).view(-1)
-    off_diag_index = torch.arange(0, diag_size - 1, dtype=torch.long, device=tensor.device)
-    off_diag_index = off_diag_index.unsqueeze(1).repeat(batch_size, 1).view(-1)
-
-    if upper:
-        chol_trid = chol_trid.transpose(-1, -2)
-
-    chol_main_diag = chol_trid[batch_index, diag_index, diag_index].view(batch_size, diag_size)
-    chol_off_diag = chol_trid[off_batch_index, off_diag_index + 1, off_diag_index].view(batch_size, diag_size - 1)
-
-    chol_solution = torch.empty_like(tensor)
-    chol_solution[:, 0, :].copy_(tensor[:, 0, :] / chol_main_diag[:, 0].unsqueeze(-1))
-    for i in range(1, diag_size):
-        inner_part = tensor[:, i, :]
-        inner_part = inner_part - chol_off_diag[:, i - 1].unsqueeze(-1) * chol_solution[:, i - 1, :]
-        chol_solution[:, i, :].copy_(inner_part / chol_main_diag[:, i].unsqueeze(-1))
-
-    solution = torch.empty_like(chol_solution)
-    solution[:, -1, :].copy_(chol_solution[:, -1, :] / chol_main_diag[:, -1].unsqueeze(-1))
-    for i in range(diag_size - 2, -1, -1):
-        inner_part = chol_solution[:, i, :] - chol_off_diag[:, i].unsqueeze(-1) * solution[:, i + 1, :]
-        solution[:, i, :].copy_(inner_part / chol_main_diag[:, i].unsqueeze(-1))
-
-    return solution
-
-
 def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
     """Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
     Args:
@@ -123,3 +42,10 @@ def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
 
     if out is None:
         return L
+
+
+def cholesky_solve(b, u, upper=False):
+    if hasattr(torch, "cholesky_solve"):
+        return torch.cholesky_solve(b, u, upper=False)
+    else:
+        return torch.potrs(b, u, upper=False)
diff --git a/gpytorch/utils/pivoted_cholesky.py b/gpytorch/utils/pivoted_cholesky.py
index 880c2bc7..9f2b6205 100644
--- a/gpytorch/utils/pivoted_cholesky.py
+++ b/gpytorch/utils/pivoted_cholesky.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import torch
+from .cholesky import cholesky_solve
 
 
 def pivoted_cholesky(matrix, max_iter, error_tol=1e-3):
@@ -98,7 +99,7 @@ def woodbury_factor(low_rank_mat, shift):
 
     shifted_mat = shifted_mat + torch.eye(k, dtype=shifted_mat.dtype, device=shifted_mat.device)
 
-    R = torch.potrs(low_rank_mat, torch.cholesky(shifted_mat, upper=True))
+    R = cholesky_solve(low_rank_mat, torch.cholesky(shifted_mat))
     return R
 
 


Test patch:
diff --git a/test/utils/test_cholesky.py b/test/utils/test_cholesky.py
index f349cd88..a0a5f370 100644
--- a/test/utils/test_cholesky.py
+++ b/test/utils/test_cholesky.py
@@ -2,26 +2,9 @@
 
 import unittest
 import warnings
-from test._utils import approx_equal
 
 import torch
-from gpytorch.utils.cholesky import psd_safe_cholesky, tridiag_batch_potrf, tridiag_batch_potrs
-
-
-class TestTriDiag(unittest.TestCase):
-    def test_potrf(self):
-        chol = torch.tensor([[1, 0, 0, 0], [2, 1, 0, 0], [0, 1, 2, 0], [0, 0, 2, 3]], dtype=torch.float).unsqueeze(0)
-        trid = chol.matmul(chol.transpose(-1, -2))
-
-        self.assertTrue(torch.equal(chol, tridiag_batch_potrf(trid, upper=False)))
-
-    def test_potrs(self):
-        chol = torch.tensor([[1, 0, 0, 0], [2, 1, 0, 0], [0, 1, 2, 0], [0, 0, 2, 3]], dtype=torch.float).unsqueeze(0)
-
-        mat = torch.randn(1, 4, 3)
-        self.assertTrue(
-            approx_equal(torch.potrs(mat[0], chol[0], upper=False), tridiag_batch_potrs(mat, chol, upper=False)[0])
-        )
+from gpytorch.utils.cholesky import psd_safe_cholesky
 
 
 class TestPSDSafeCholesky(unittest.TestCase):
diff --git a/test/utils/test_linear_cg.py b/test/utils/test_linear_cg.py
index 193b25cf..2fa064b5 100644
--- a/test/utils/test_linear_cg.py
+++ b/test/utils/test_linear_cg.py
@@ -6,6 +6,7 @@ import torch
 import unittest
 from test._utils import approx_equal
 from gpytorch.utils.linear_cg import linear_cg
+from gpytorch.utils.cholesky import cholesky_solve
 
 
 class TestLinearCG(unittest.TestCase):
@@ -32,8 +33,8 @@ class TestLinearCG(unittest.TestCase):
         solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
 
         # Check cg
-        matrix_chol = matrix.cholesky(upper=True)
-        actual = torch.potrs(rhs, matrix_chol)
+        matrix_chol = matrix.cholesky()
+        actual = cholesky_solve(rhs, matrix_chol)
         self.assertTrue(approx_equal(solves, actual))
 
     def test_cg_with_tridiag(self):
@@ -49,8 +50,8 @@ class TestLinearCG(unittest.TestCase):
         )
 
         # Check cg
-        matrix_chol = matrix.cholesky(upper=True)
-        actual = torch.potrs(rhs, matrix_chol)
+        matrix_chol = matrix.cholesky()
+        actual = cholesky_solve(rhs, matrix_chol)
         self.assertTrue(approx_equal(solves, actual))
 
         # Check tridiag
@@ -71,8 +72,8 @@ class TestLinearCG(unittest.TestCase):
         solves = linear_cg(matrix.matmul, rhs=rhs, max_iter=size)
 
         # Check cg
-        matrix_chol = torch.cholesky(matrix, upper=True)
-        actual = torch.potrs(rhs, matrix_chol)
+        matrix_chol = torch.cholesky(matrix)
+        actual = cholesky_solve(rhs, matrix_chol)
         self.assertTrue(approx_equal(solves, actual))
 
     def test_batch_cg_with_tridiag(self):
@@ -89,8 +90,8 @@ class TestLinearCG(unittest.TestCase):
         )
 
         # Check cg
-        matrix_chol = torch.cholesky(matrix, upper=True)
-        actual = torch.potrs(rhs, matrix_chol)
+        matrix_chol = torch.cholesky(matrix)
+        actual = cholesky_solve(rhs, matrix_chol)
         self.assertTrue(approx_equal(solves, actual))
 
         # Check tridiag