Repository: cornellius-gp/gpytorch
Pull request #141 (2018-05-23 22:56:16)

Lazily evaluate kernel functions
Working on fixing #135. So far, simple GPs and the 1D KISS-GP tests pass.

My initial solution to this is to make the `__call__` method of `Kernel` directly create the new `LazyEvaluatedKernelVariable`, which basically supports a `__getitem__` method (so we can subindex in to a kernel matrix before evaluating the kernel, which is the point of this PR) and an `evaluate()` method. This is a meta lazy variable, in that the `evaluate()` method can return another `LazyVariable` if the kernel call does.

The reason I am obviously unhappy with this is that if the user does anything in forward other than call a kernel and return the `LazyEvaluatedKernelVariable`, we now call `evaluate()` in the model, which is obviously a Bad Thing if the covariance is, for example, already an `InterpolatedLazyVariable` or something.

I have a decent solution to this that I'll work on either tomorrow or after graduation when I have time.

(P.S.  I think the "obvious" solution of just refactoring so that we call forward on the different sets of data directly would actually require an even larger internals rework because of how we handle predictions via `LazyVariable` methods.) 

jacobrgardner: Basically, what I'm thinking is that the new object isn't a LazyVariable at all, and just "evaluates" the kernel directly in the `__getitem__` method. I was originally not doing this because I was going to have it define a diag() method so that we could get variances more easily for exact GPs. However after working on it for a bit I think a better solution is just to define an `LazyVariable_exact_predictive_var(...)` that by default does `self._exact_predictive_covar(...).diag()`
gpleiss: It might make sense to have this object not be a LazyVariable. We could create a new class of Lazy objects.

GPyTorch... the world's laziest library 🤣 
jacobrgardner: GPyTorch follows the __aggressively lazy__ design pattern. We refuse to actually compute anything until you actually ask us at least twice :-)
jacobrgardner: Ok, this PR is essentially done now. It fixes #135 and should make @Balandat and @darbour's solution to #113 work.

## Internal changes
- The `Kernel.__call__` method now returns a new `LazyEvaluatedKernelVariable` rather than calling forward directly.
- `LazyEvaluatedKernelVariable` supports sub indexing, and returns a new `LazyEvaluatedKernelVariable` from `__getitem__`. It defines an `evaluate_kernel` method that actually calls the kernel. Thus, kernel calls happen after subindexing, not before. 
- `LazyEvaluatedKernelVariable` overwrites the `representation` and `representation_tree` functions to masquerade as the `LazyVariable` that results from calling `evaluate_kernel` rather than as itself. Reconstructing a `LazyVariable` from its representation tree therefore evaluates all kernel calls. Because all functions that operate on `LazyVariables` reconstruct it, these functions will work even with `LazyEvaluatedKernelVariables` in the tree.
- __For a special case__, the base `LazyVariable` class now defines an `evaluate_kernel` method that just rebuilds self. This is because `pivoted_cholesky` is not a function yet, so it (and future things like it) may need a quick way to do this.

## User level changes
- All kernels are now lazily evaluated. Therefore, even basic kernels like `RBFKernel` need to be followed with an evaluate() call if a tensor is desired. For example, `RBFKernel(x1, x2).evaluate()`.
- Nothing changes at the actual "using a GP" level, because all of this is handled in `exact_predictive_mean` and `exact_predictive_covar`.
- Nothing changes for sparse GP methods, because those kernels return `LazyVariables` that would need to be evaluated to get tensors anyways.

## Future work
- Once `pivoted_cholesky` is a `Function`, the last remaining calls to `LazyVariable.evaluate_kernel` can be removed from the code base. This is a PR in its own right, however.
- `LazyEvaluatedKernelVariable` defines a `diag` function that right now just calls `self.evaluate_kernel().diag()`, but could do this more intelligently using batch indexing. This should allow us to somewhat easily define a "diagonal only" prediction mode that gets variances instead of covariances. This is not necessarily useful for the sparse GP methods when using fast predictive variances, but is useful for exact GPs.



Patch:
diff --git a/gpytorch/kernels/grid_kernel.py b/gpytorch/kernels/grid_kernel.py
index 550a1f0b..a94c8f45 100644
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -6,7 +6,7 @@ from __future__ import unicode_literals
 import torch
 from torch.autograd import Variable
 from .kernel import Kernel
-from ..lazy import ToeplitzLazyVariable, KroneckerProductLazyVariable, NonLazyVariable
+from ..lazy import ToeplitzLazyVariable, KroneckerProductLazyVariable
 from .. import settings
 
 
@@ -43,15 +43,15 @@ class GridKernel(Kernel):
 
             if settings.use_toeplitz.on():
                 first_item = grid_var[:, 0:1].contiguous()
-                covar_columns = self.base_kernel_module(first_item, grid_var, **kwargs)
+                covar_columns = self.base_kernel_module(first_item, grid_var, **kwargs).evaluate()
                 covars = [
                     ToeplitzLazyVariable(covar_columns[i : i + 1].squeeze(-2))
                     for i in range(n_dim)
                 ]
             else:
                 grid_var = grid_var.view(n_dim, -1, 1)
-                covars = self.base_kernel_module(grid_var, grid_var, **kwargs)
-                covars = [NonLazyVariable(covars[i : i + 1]) for i in range(n_dim)]
+                covars = self.base_kernel_module(grid_var, grid_var, **kwargs).evaluate_kernel()
+                covars = [covars[i : i + 1] for i in range(n_dim)]
 
             if n_dim > 1:
                 covar = KroneckerProductLazyVariable(*covars)
diff --git a/gpytorch/kernels/inducing_point_kernel.py b/gpytorch/kernels/inducing_point_kernel.py
index e2f730d7..bf1f554f 100644
--- a/gpytorch/kernels/inducing_point_kernel.py
+++ b/gpytorch/kernels/inducing_point_kernel.py
@@ -38,7 +38,7 @@ class InducingPointKernel(Kernel):
         if not self.training and hasattr(self, "_cached_kernel_mat"):
             return self._cached_kernel_mat
         else:
-            res = self.base_kernel_module(self.inducing_points, self.inducing_points)
+            res = self.base_kernel_module(self.inducing_points, self.inducing_points).evaluate()
             if not self.training:
                 self._cached_kernel_mat = res
             return res
@@ -61,11 +61,11 @@ class InducingPointKernel(Kernel):
             return res
 
     def _get_covariance(self, x1, x2):
-        k_ux1 = self.base_kernel_module(x1, self.inducing_points)
+        k_ux1 = self.base_kernel_module(x1, self.inducing_points).evaluate()
         if torch.equal(x1, x2):
             covar = RootLazyVariable(k_ux1.matmul(self._inducing_inv_root))
         else:
-            k_ux2 = self.base_kernel_module(x2, self.inducing_points)
+            k_ux2 = self.base_kernel_module(x2, self.inducing_points).evaluate()
             covar = MatmulLazyVariable(
                 k_ux1.matmul(self._inducing_inv_root),
                 k_ux2.matmul(self._inducing_inv_root).transpose(-1, -2),
diff --git a/gpytorch/kernels/kernel.py b/gpytorch/kernels/kernel.py
index f15cbe65..77640be2 100644
--- a/gpytorch/kernels/kernel.py
+++ b/gpytorch/kernels/kernel.py
@@ -5,6 +5,7 @@ from __future__ import unicode_literals
 
 import torch
 from ..module import Module
+from ..lazy import LazyEvaluatedKernelVariable
 
 
 class Kernel(Module):
@@ -57,16 +58,7 @@ class Kernel(Module):
         if not x1.size(-1) == x2.size(-1):
             raise RuntimeError("x1 and x2 must have the same number of dimensions!")
 
-        # Do everything in batch mode by default
-        is_batch = x1.ndimension() == 3
-        if not is_batch:
-            x1 = x1.unsqueeze(0)
-            x2 = x2.unsqueeze(0)
-
-        res = super(Kernel, self).__call__(x1, x2, **params)
-        if not is_batch:
-            res = res[0]
-        return res
+        return LazyEvaluatedKernelVariable(self, x1, x2)
 
     def __add__(self, other):
         return AdditiveKernel(self, other)
diff --git a/gpytorch/kernels/multiplicative_grid_interpolation_kernel.py b/gpytorch/kernels/multiplicative_grid_interpolation_kernel.py
index 4a3454f3..3bb94383 100644
--- a/gpytorch/kernels/multiplicative_grid_interpolation_kernel.py
+++ b/gpytorch/kernels/multiplicative_grid_interpolation_kernel.py
@@ -44,3 +44,15 @@ class MultiplicativeGridInterpolationKernel(GridInterpolationKernel):
         res = super(MultiplicativeGridInterpolationKernel, self).forward(x1, x2)
         res = res.mul_batch(mul_batch_size=self.n_components)
         return res
+
+    def __call__(self, x1_, x2_=None, **params):
+        """
+        We cannot lazily evaluate actual kernel calls when using SKIP, because we cannot root decompose
+        rectangular matrices.
+
+        Because we slice in to the kernel during prediction to get the test x train covar before
+        calling evaluate_kernel, the order of operations would mean we would get a MulLazyVariable representing a
+        rectangular matrix, which we cannot matmul with because we cannot root decompose it. Thus, SKIP actually
+        *requires* that we work with the full (train + test) x (train + test) kernel matrix.
+        """
+        return super(MultiplicativeGridInterpolationKernel, self).__call__(x1_, x2_, **params).evaluate_kernel()
diff --git a/gpytorch/lazy/__init__.py b/gpytorch/lazy/__init__.py
index e143061c..af40b832 100644
--- a/gpytorch/lazy/__init__.py
+++ b/gpytorch/lazy/__init__.py
@@ -6,6 +6,7 @@ from .constant_mul_lazy_variable import ConstantMulLazyVariable
 from .diag_lazy_variable import DiagLazyVariable
 from .interpolated_lazy_variable import InterpolatedLazyVariable
 from .kronecker_product_lazy_variable import KroneckerProductLazyVariable
+from .lazy_evaluated_kernel_variable import LazyEvaluatedKernelVariable
 from .matmul_lazy_variable import MatmulLazyVariable
 from .mul_lazy_variable import MulLazyVariable
 from .non_lazy_variable import NonLazyVariable
@@ -25,6 +26,7 @@ __all__ = [
     DiagLazyVariable,
     InterpolatedLazyVariable,
     KroneckerProductLazyVariable,
+    LazyEvaluatedKernelVariable,
     MatmulLazyVariable,
     MulLazyVariable,
     NonLazyVariable,
diff --git a/gpytorch/lazy/constant_mul_lazy_variable.py b/gpytorch/lazy/constant_mul_lazy_variable.py
index c2c567b8..e951424e 100644
--- a/gpytorch/lazy/constant_mul_lazy_variable.py
+++ b/gpytorch/lazy/constant_mul_lazy_variable.py
@@ -59,6 +59,10 @@ class ConstantMulLazyVariable(LazyVariable):
         res = self.lazy_var._approx_diag()
         return res * self.constant.expand_as(res)
 
+    def evaluate(self):
+        res = self.lazy_var.evaluate()
+        return res * self.constant.expand_as(res)
+
     def diag(self):
         res = self.lazy_var.diag()
         res = res * self.constant.expand_as(res)
diff --git a/gpytorch/lazy/lazy_evaluated_kernel_variable.py b/gpytorch/lazy/lazy_evaluated_kernel_variable.py
new file mode 100644
index 00000000..e3f29501
--- /dev/null
+++ b/gpytorch/lazy/lazy_evaluated_kernel_variable.py
@@ -0,0 +1,93 @@
+from .lazy_variable import LazyVariable
+from .non_lazy_variable import NonLazyVariable
+from .lazy_variable_representation_tree import LazyVariableRepresentationTree
+import torch
+
+
+class LazyEvaluatedKernelVariable(LazyVariable):
+    def __init__(self, kernel, x1, x2, **params):
+        super(LazyEvaluatedKernelVariable, self).__init__(kernel, x1, x2, **params)
+        self.kernel = kernel
+        self.x1 = x1
+        self.x2 = x2
+        self.is_batch = self.x1.ndimension() == 3
+
+    def _matmul(self, rhs):
+        raise RuntimeError('A LazyEvaluatedKernelVariable is not intended to be used directly as a tensor!'
+                           ' Call evaluate() first.')
+
+    def _t_matmul(self, rhs):
+        raise RuntimeError('A LazyEvaluatedKernelVariable is not intended to be used directly as a tensor!'
+                           ' Call evaluate() first.')
+
+    def _quad_form_derivative(self, left_vecs, right_vecs):
+        raise RuntimeError('A LazyEvaluatedKernelVariable is not intended to be used directly as a tensor!'
+                           ' Call evaluate() first.')
+
+    def diag(self):
+        """
+        TODO: Can be handled by calling the kernel after creating some new batch dimensions and transposing.
+
+        Hopefully, this can be easily used to add a 'variance_only' prediction mode.
+        """
+        return self.evaluate_kernel().diag()
+
+    def evaluate_kernel(self):
+        """
+        NB: This is a meta LazyVariable, in the sense that evaluate can return a LazyVariable if the kernel being
+        evaluated does so.
+        """
+        from ..kernels import Kernel
+        if hasattr(self, '_cached_kernel_eval'):
+            return self._cached_kernel_eval
+
+        else:
+            if not self.is_batch:
+                x1 = self.x1.unsqueeze(0)
+                x2 = self.x2.unsqueeze(0)
+            else:
+                x1 = self.x1
+                x2 = self.x2
+
+            self._cached_kernel_eval = super(Kernel, self.kernel).__call__(x1, x2)
+
+            if not self.is_batch:
+                self._cached_kernel_eval = self._cached_kernel_eval[0]
+
+            if not isinstance(self._cached_kernel_eval, LazyVariable):
+                self._cached_kernel_eval = NonLazyVariable(self._cached_kernel_eval)
+
+            return self._cached_kernel_eval
+
+    def representation(self):
+        return self.evaluate_kernel().representation()
+
+    def representation_tree(self):
+        return LazyVariableRepresentationTree(self.evaluate_kernel())
+
+    def evaluate(self):
+        return self.evaluate_kernel().evaluate()
+
+    def __getitem__(self, index):
+        index = list(index) if isinstance(index, tuple) else [index]
+        ndimension = self.ndimension()
+        index += [slice(None, None, None)] * (ndimension - len(index))
+        if self.is_batch:
+            batch_index = index[0]
+            left_index = index[1]
+            right_index = index[2]
+            return LazyEvaluatedKernelVariable(self.kernel,
+                                               self.x1[batch_index, left_index, :],
+                                               self.x2[batch_index, right_index, :])
+        else:
+            left_index = index[0]
+            right_index = index[1]
+            return LazyEvaluatedKernelVariable(self.kernel,
+                                               self.x1[left_index, :],
+                                               self.x2[right_index, :])
+
+    def _size(self):
+        if self.is_batch:
+            return torch.Size((self.x1.size(0), self.x1.size(-2), self.x2.size(-2)))
+        else:
+            return torch.Size((self.x1.size(-2), self.x2.size(-2)))
diff --git a/gpytorch/lazy/lazy_variable.py b/gpytorch/lazy/lazy_variable.py
index fad78af9..75f2847a 100644
--- a/gpytorch/lazy/lazy_variable.py
+++ b/gpytorch/lazy/lazy_variable.py
@@ -290,6 +290,13 @@ class LazyVariable(object):
                 eye = eye.unsqueeze(0).expand(batch_size, n_cols, n_cols)
             return self.matmul(eye)
 
+    def evaluate_kernel(self):
+        """
+        Return a new LazyVariable representing the same one as this one, but with all lazily evaluated kernels
+        actually evaluated.
+        """
+        return self.representation_tree()(*self.representation())
+
     def exact_predictive_mean(
         self, full_mean, train_labels, noise, precomputed_cache=None
     ):
diff --git a/gpytorch/lazy/lazy_variable_representation_tree.py b/gpytorch/lazy/lazy_variable_representation_tree.py
index dd30da10..6b74a55f 100644
--- a/gpytorch/lazy/lazy_variable_representation_tree.py
+++ b/gpytorch/lazy/lazy_variable_representation_tree.py
@@ -18,7 +18,7 @@ class LazyVariableRepresentationTree(object):
                 self.children.append(
                     (
                         slice(counter, counter + representation_size, None),
-                        LazyVariableRepresentationTree(arg),
+                        arg.representation_tree(),
                     )
                 )
                 counter += representation_size
diff --git a/gpytorch/models/abstract_variational_gp.py b/gpytorch/models/abstract_variational_gp.py
index 58bde5e9..cfe56563 100644
--- a/gpytorch/models/abstract_variational_gp.py
+++ b/gpytorch/models/abstract_variational_gp.py
@@ -75,6 +75,8 @@ class AbstractVariationalGP(Module):
                 "%s.forward must return a GaussianRandomVariable"
                 % self.__class__.__name__
             )
+
+        res = GaussianRandomVariable(res.mean(), res.covar().evaluate_kernel())
         return res
 
     def variational_output(self):
diff --git a/gpytorch/utils/pivoted_cholesky.py b/gpytorch/utils/pivoted_cholesky.py
index ea8a67a7..59c602c7 100644
--- a/gpytorch/utils/pivoted_cholesky.py
+++ b/gpytorch/utils/pivoted_cholesky.py
@@ -18,6 +18,7 @@ def pivoted_cholesky(matrix, max_iter, error_tol=1e-3):
     # Need to get diagonals. This is easy if it's a LazyVariable, since
     # LazyVariable.diag() operates in batch mode.
     if isinstance(matrix, LazyVariable):
+        matrix = matrix.evaluate_kernel()
         matrix_diag = matrix._approx_diag()
     elif isinstance(matrix, Variable):
         matrix_diag = NonLazyVariable(matrix).diag()


Test patch:
diff --git a/test/kernels/test_additive_kernel.py b/test/kernels/test_additive_kernel.py
index 0b15b7fd..3637083b 100644
--- a/test/kernels/test_additive_kernel.py
+++ b/test/kernels/test_additive_kernel.py
@@ -25,7 +25,7 @@ class TestAdditiveKernel(unittest.TestCase):
         ).exp() ** 2
 
         kernel.eval()
-        res = kernel(a, b)
+        res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 2e-5)
 
     def test_computes_sum_of_radial_basis_function(self):
@@ -42,7 +42,7 @@ class TestAdditiveKernel(unittest.TestCase):
         ).exp() * 2
 
         kernel.eval()
-        res = kernel(a, b)
+        res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 2e-5)
 
     def test_computes_sum_radial_basis_function_gradient(self):
@@ -62,7 +62,7 @@ class TestAdditiveKernel(unittest.TestCase):
         kernel = kernel_1 + kernel_2
         kernel.eval()
 
-        output = kernel(a, b)
+        output = kernel(a, b).evaluate()
         output.backward(gradient=torch.eye(3))
         res = (
             kernel.kernel_1.log_lengthscale.grad + kernel.kernel_2.log_lengthscale.grad
diff --git a/test/kernels/test_periodic_kernel.py b/test/kernels/test_periodic_kernel.py
index 1825198c..ff123c2b 100644
--- a/test/kernels/test_periodic_kernel.py
+++ b/test/kernels/test_periodic_kernel.py
@@ -28,7 +28,7 @@ class TestPeriodicKernel(unittest.TestCase):
                 val = 2 * torch.pow(torch.sin(math.pi * (a[i] - b[j])), 2) / lengthscale
                 actual[i, j] = torch.exp(-val)[0]
 
-        res = kernel(Variable(a), Variable(b)).data
+        res = kernel(Variable(a), Variable(b)).evaluate().data
         self.assertLess(torch.norm(res - actual), 1e-5)
 
 
diff --git a/test/kernels/test_rbf_kernel.py b/test/kernels/test_rbf_kernel.py
index 34f1daf8..b23553fe 100644
--- a/test/kernels/test_rbf_kernel.py
+++ b/test/kernels/test_rbf_kernel.py
@@ -21,7 +21,7 @@ class TestRBFKernel(unittest.TestCase):
         kernel.eval()
 
         actual = (a - b).div_(lengthscales).pow(2).sum(dim=-1).mul_(-0.5).exp()
-        res = kernel(a, b)
+        res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual.unsqueeze(-1)), 1e-5)
 
     def test_subset_active_compute_radial_basis_function(self):
@@ -38,7 +38,7 @@ class TestRBFKernel(unittest.TestCase):
         actual = torch.Tensor([[16, 4], [4, 0], [64, 36]]).mul_(-0.5).div_(
             lengthscale ** 2
         ).exp()
-        res = kernel(a, b)
+        res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_computes_radial_basis_function(self):
@@ -52,7 +52,7 @@ class TestRBFKernel(unittest.TestCase):
         actual = torch.Tensor([[16, 4], [4, 0], [64, 36]]).mul_(-0.5).div_(
             lengthscale ** 2
         ).exp()
-        res = kernel(a, b)
+        res = kernel(a, b).evaluate()
         self.assertLess(torch.norm(res - actual), 1e-5)
 
     def test_computes_radial_basis_function_gradient(self):
@@ -70,7 +70,7 @@ class TestRBFKernel(unittest.TestCase):
         actual_output.backward(gradient=torch.eye(3))
         actual_param_grad = param.grad.sum()
 
-        output = kernel(a, b)
+        output = kernel(a, b).evaluate()
         output.backward(gradient=torch.eye(3))
         res = kernel.log_lengthscale.grad
 
@@ -93,7 +93,7 @@ class TestRBFKernel(unittest.TestCase):
         kernel = RBFKernel(active_dims=[0])
         kernel.initialize(log_lengthscale=math.log(lengthscale))
         kernel.eval()
-        output = kernel(a, b)
+        output = kernel(a, b).evaluate()
         output.backward(gradient=torch.eye(3))
         res = kernel.log_lengthscale.grad
 
diff --git a/test/util/test_pivoted_cholesky.py b/test/util/test_pivoted_cholesky.py
index 7a8f755c..d2bc03e4 100644
--- a/test/util/test_pivoted_cholesky.py
+++ b/test/util/test_pivoted_cholesky.py
@@ -14,7 +14,7 @@ class TestPivotedCholesky(unittest.TestCase):
     def test_pivoted_cholesky(self):
         size = 100
         train_x = torch.linspace(0, 1, size)
-        covar_matrix = RBFKernel()(train_x, train_x)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
         covar_approx = piv_chol.t().matmul(piv_chol)
         self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
@@ -22,7 +22,7 @@ class TestPivotedCholesky(unittest.TestCase):
     def test_solve_vector(self):
         size = 100
         train_x = torch.linspace(0, 1, size)
-        covar_matrix = RBFKernel()(train_x, train_x)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
         woodbury_factor = pivoted_cholesky.woodbury_factor(
             piv_chol, torch.Tensor(100).fill_(1)
@@ -40,7 +40,7 @@ class TestPivotedCholesky(unittest.TestCase):
     def test_solve(self):
         size = 100
         train_x = torch.linspace(0, 1, size)
-        covar_matrix = RBFKernel()(train_x, train_x)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
         woodbury_factor = pivoted_cholesky.woodbury_factor(
             piv_chol, torch.Tensor(100).fill_(1)
@@ -69,7 +69,7 @@ class TestPivotedCholeskyBatch(unittest.TestCase):
         ).unsqueeze(
             -1
         )
-        covar_matrix = RBFKernel()(train_x, train_x)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
         covar_approx = piv_chol.transpose(1, 2).matmul(piv_chol)
 
@@ -86,7 +86,7 @@ class TestPivotedCholeskyBatch(unittest.TestCase):
         ).unsqueeze(
             -1
         )
-        covar_matrix = RBFKernel()(train_x, train_x)
+        covar_matrix = RBFKernel()(train_x, train_x).evaluate()
         piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
         woodbury_factor = pivoted_cholesky.woodbury_factor(
             piv_chol, torch.Tensor(2, 100).fill_(1)