Implementing linear conjugate gradients

[kunalghosh]

This is a draft implementation of the linear conjugate gradients algorithm.
It is designed to be output the same as gpytorch.util.linear_cg and can be tested to give the same results using the code snippet below.

Let's first generate the data:

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)
print(np.diag(K))

y = np.random.randn(N, 1)
y_torch = torch.from_numpy(y)

Now we get the output from GPyTorch

result, tridiag = gpytorch.utils.linear_cg(K_torch, y_torch, n_tridiag=1)

And also from our implementation:

import pymc_experimental as pymx
result_pmx, tridiag_pmx = pymx.utils.linear_cg(K, y, n_tridiag=1)

Check the outputs

assert np.allclose(result, result_pmx) is True, "BUG: result doesn't match gpytorch values"
assert np.allclose(tridiag tridiag_pmx) is True, "BUG: tridiagonal values don't match gpytorch values"

[kunalghosh]

@bwengals Ping, can you take a look at the PR please :)

[bwengals]

will do. It might take me a couple days, it's a lot to go through, but it looks like you got it working, congrats!

[kunalghosh]

Just a heads up, I noticed this check fail yesterday. Was passing when I ran it on Colab, looking into it.

[kunalghosh]

Fixed it should have been if not np.allclose(...)

[bwengals]

Since there are only two parameters here, could you pass eval_cg_tolerance and cg_tolerance as args to linear_cg? I think it would simplify things a little bit, since you wouldn't need this class.

[kunalghosh]

This class is just a place holder, I would ideally put these two values in some sort of default settings file or class used by PyMC. I moved the constants to be global variables and kept the logic same as in GPyTorch.

[bwengals]

There's no need to follow GPyTorch's particular way of doing things -- in PyMC there aren't any global variables or settings files, that I'm aware of. It's important to be consistent with these sort of things across the codebase.

[bwengals]

This is just a small nitpick, but you could combine this, so you don't need to create the linear_cg_retvals.

[kunalghosh]

Fixed in the latest commit

[bwengals]

It would be nice if you included your example code to test this out in the actual PR somewhere! You could make a brief notebook that demonstrates how to use it.

I left some minor comments, but I think shouldn't hold up the merge process. Depending on the strategy, this implementation will be translated to aesara/numba/jax, so the goal of this code should be clarity. Feel free to make changes and merge it when you're ready.

[bwengals]

BTW is this still PR still in draft mode, or do you think you're pretty much done?

[kunalghosh]

BTW is this still PR still in draft mode, or do you think you're pretty much done?

Nope, this is ready to check-in now :)

[bwengals]

Sorry for the delay, everything looks copacetic