Initial commit of pivoted Cholesky algorithm from GPyTorch

[kunalghosh]

What is this PR about?

Fast exact Gaussian processes (Gardner et.al 2018) works on a modified batch conjugate gradient descent algorithm which needs a function to return the preconditioning matrix, given a kernel. This PR implements the function to compute the preconditioning matrix.

This produces the same results as gpytorch.pivoted_cholesky()

Test kernel matrices can be generated using the following code which generates random positive semi-definite matrices

import torch 
import gpytorch
import numpy as np

# Test PSD Matrix

N = 10
rank = 5
np.random.seed(1234) # nans with seed 1234
K = np.random.randn(N, N)
K = K @ K.T + N * np.eye(N)
K_torch = torch.from_numpy(K)

Now passing the same kernel matrix to the two functions will yield the same output.

import pymc_experimental as pymx

L_gpt = gpytorch.pivoted_cholesky(K_torch, rank=rank, error_tol=1e-3)
L_np, _  = pymx.utils.pivoted_cholesky(K, max_iter=rank, error_tol=1e-3)
assert np.allclose(L_gpt, L_np.T) == True, "BUG: The two cholesky decompositions are not same !"

[kunalghosh]

@bwengals I have copied this PR over from pymc

[bwengals]

I think It'd be best to make torch and gpytorch optional imports.

try:
    import torch
    import gpytorch
except ImportError as e:
    raise ImportError("message")

That way they don't need to be dependencies for someone who wan'ts to use pymc_experimental for something else.

[aloctavodia]

Agreed

[ferrine]

I've looked into the apply_permutation function, I think it is better to avoid the dependency and implement the same function here

[kunalghosh]

@ferrine Yes, I plan to use the gpytorch.utils.permutation apply_permutation temporarily and implement the apply_permutation function once I have GP inference working on the GPU.

I have now added the try...catch checks. Is it ok if I check-in this ? I have created a n issue for the same #85

[bwengals]

100% agree @ferrine, longer term that'll have to implemented here, but this PRs more at the proof-of-concept level so lot's of time to go back over those details later. With the try...except it shouldn't get in anyones way.

[bwengals]

Same comments as the other PR, ready to be taken off of draft?

[aloctavodia]

Why you comment the test?

[kunalghosh]

This was triggering the import of GPyTorch and Torch which would cause the tests to fail. I thought it would be best to comment out the tests until I remove the dependency of my code on GPyTorch and Torch.

[aloctavodia]

instead of a print use raise ImportError(...)

[bwengals]

definitely, also consider raise ImportError(...) from e https://stackoverflow.com/questions/24752395/python-raise-from-usage

[aloctavodia]

me being annoying: np.inf

[bwengals]

Looks great @kunalghosh! Feel free to hit the green button whenever you're ready

[bwengals]

definitely, also consider raise ImportError(...) from e https://stackoverflow.com/questions/24752395/python-raise-from-usage

[bwengals]

Sorry, didn't notice the conflict on pymc_experimental/utils/__init__.py. You'll have to resolve that before merging.

[bwengals]

nice job @kunalghosh, sorry for the delay on my side. This is good progress.