def test_root_inv_decomposition(self):
# Forward
probe_vectors = torch.randn(4, 5)
test_vectors = torch.randn(4, 5)
root = NonLazyTensor(self.mat).root_inv_decomposition(
probe_vectors, test_vectors).root.evaluate()
res = root.matmul(root.transpose(-1, -2))
actual = self.mat_clone.inverse()
self.assertTrue(approx_equal(res, actual))
# Backward
res.trace().backward()
actual.trace().backward()
self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
def test_matmul_multiple_vecs(self):
# Forward
res = NonLazyTensor(self.mat_var).matmul(self.vecs_var)
actual = self.mat_var_copy.matmul(self.vecs_var_copy)
self.assertTrue(approx_equal(res, actual))
# Backward
grad_output = torch.randn(3, 4)
res.backward(gradient=grad_output)
actual.backward(gradient=grad_output)
self.assertTrue(approx_equal(self.mat_var_copy.grad,
self.mat_var.grad))
self.assertTrue(
approx_equal(self.vecs_var_copy.grad, self.vecs_var.grad))
def test_root_decomposition(self):
# Forward
root = NonLazyTensor(self.mat).root_decomposition().root.evaluate()
res = root.matmul(root.transpose(-1, -2))
self.assertTrue(approx_equal(res, self.mat))
# Backward
sum([mat.trace()
for mat in res.view(-1, *self.mat.shape[-2:])]).backward()
sum([
mat.trace()
for mat in self.mat_clone.view(-1, *self.mat.shape[-2:])
]).backward()
self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
def test_lkj_prior_batch_log_prob(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
prior = LKJPrior(2, torch.tensor([0.5, 1.5], device=device))
S = torch.eye(2, device=device)
self.assertTrue(
approx_equal(prior.log_prob(S),
torch.tensor([-1.86942, -0.483129], device=S.device)))
S = torch.stack(
[S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
self.assertTrue(
approx_equal(prior.log_prob(S),
torch.tensor([-1.86942, -0.62697], device=S.device)))
with self.assertRaises(ValueError):
prior.log_prob(torch.eye(3, device=device))
def test_get_item_tensor_index(self):
# Tests the default LV.__getitem__ behavior
lazy_tensor = ZeroLazyTensor(5, 5)
evaluated = lazy_tensor.evaluate()
index = (torch.tensor([0, 0, 1, 2]), torch.tensor([0, 1, 0, 2]))
self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
index = (torch.tensor([0, 0, 1, 2]), slice(None, None, None))
self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
index = (slice(None, None, None), torch.tensor([0, 0, 1, 2]))
self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
index = (Ellipsis, slice(None, None, None), torch.tensor([0, 0, 1, 2]))
self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
index = (Ellipsis, torch.tensor([0, 0, 1, 2]))
self.assertTrue(approx_equal(lazy_tensor[index], evaluated[index]))
def test_root_inv_decomposition(self):
# Forward
probe_vectors = torch.randn(3, 4, 5)
test_vectors = torch.randn(3, 4, 5)
root = NonLazyTensor(self.mat).root_inv_decomposition(
probe_vectors, test_vectors).root.evaluate()
res = root.matmul(root.transpose(-1, -2))
actual = torch.cat(
[mat.inverse().unsqueeze(0) for mat in self.mat_clone])
self.assertTrue(approx_equal(res, actual))
# Backward
sum([mat.trace() for mat in res]).backward()
sum([mat.trace() for mat in actual]).backward()
self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
def test_fantasy_updates_batch(self, cuda=False):
train_x, test_x, train_y, test_y = self._get_data(cuda=cuda)
# We're manually going to set the hyperparameters to something they shouldn't be
likelihood = GaussianLikelihood()
gp_model = ExactGPModel(train_x, train_y, likelihood)
mll = gpytorch.ExactMarginalLogLikelihood(likelihood, gp_model)
gp_model.covar_module.base_kernel.initialize(lengthscale=exp(1))
gp_model.mean_module.initialize(constant=0)
likelihood.initialize(noise=exp(1))
if cuda:
gp_model.cuda()
likelihood.cuda()
# Find optimal model hyperparameters
gp_model.train()
likelihood.train()
optimizer = optim.Adam(list(gp_model.parameters()) +
list(likelihood.parameters()),
lr=0.15)
optimizer.n_iter = 0
for _ in range(50):
optimizer.zero_grad()
with gpytorch.settings.debug(False):
output = gp_model(train_x)
loss = -mll(output, train_y)
loss.backward()
optimizer.n_iter += 1
optimizer.step()
for param in gp_model.parameters():
self.assertTrue(param.grad is not None)
self.assertGreater(param.grad.norm().item(), 0)
for param in likelihood.parameters():
self.assertTrue(param.grad is not None)
self.assertGreater(param.grad.norm().item(), 0)
optimizer.step()
with gpytorch.settings.fast_pred_var():
# Test the model
gp_model.eval()
likelihood.eval()
test_function_predictions = likelihood(gp_model(test_x))
# Cut data down, and then add back via the fantasy interface
gp_model.set_train_data(train_x[:5], train_y[:5], strict=False)
likelihood(gp_model(test_x))
fantasy_x = train_x[5:].clone().unsqueeze(0).unsqueeze(-1).repeat(
3, 1, 1).requires_grad_(True)
fantasy_y = train_y[5:].unsqueeze(0).repeat(3, 1)
fant_model = gp_model.get_fantasy_model(fantasy_x, fantasy_y)
fant_function_predictions = likelihood(fant_model(test_x))
self.assertTrue(
approx_equal(test_function_predictions.mean,
fant_function_predictions.mean[0]))
fant_function_predictions.mean.sum().backward()
self.assertTrue(fantasy_x.grad is not None)
def test_pivoted_cholesky(self):
size = 100
train_x = torch.linspace(0, 1, size)
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))
def test_solve(self):
size = 100
train_x = torch.cat([
torch.linspace(0, 1, size).unsqueeze(0),
torch.linspace(0, 0.5, size).unsqueeze(0)
], 0).unsqueeze(-1)
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.ones(2, 100))
rhs_vector = torch.randn(2, 100, 5)
shifted_covar_matrix = covar_matrix + torch.eye(size)
real_solve = torch.cat(
[
shifted_covar_matrix[0].inverse().matmul(
rhs_vector[0]).unsqueeze(0),
shifted_covar_matrix[1].inverse().matmul(
rhs_vector[1]).unsqueeze(0),
],
0,
)
approx_solve = pivoted_cholesky.woodbury_solve(rhs_vector, piv_chol,
woodbury_factor,
torch.ones(2, 100))
self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
def test_batch_cg_with_tridiag(self):
batch = 5
size = 10
matrix = torch.randn(batch, size, size, dtype=torch.float64)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))
rhs = torch.randn(batch, size, 50, dtype=torch.float64)
solves, t_mats = linear_cg(matrix.matmul,
rhs=rhs,
n_tridiag=8,
max_iter=size,
max_tridiag_iter=10,
tolerance=0,
eps=1e-20)
# Check cg
matrix_chol = torch.cholesky(matrix, upper=True)
actual = torch.potrs(rhs, matrix_chol)
self.assertTrue(approx_equal(solves, actual))
# Check tridiag
for i in range(5):
eigs = matrix[i].symeig()[0]
for j in range(8):
approx_eigs = t_mats[j, i].symeig()[0]
self.assertLess(
torch.mean(torch.abs((eigs - approx_eigs) / eigs)), 0.05)
def test_pivoted_cholesky(self):
size = 100
train_x = torch.cat(
[
torch.linspace(0, 1, size).unsqueeze(0),
torch.linspace(0, 0.5, size).unsqueeze(0),
torch.linspace(0, 0.25, size).unsqueeze(0),
torch.linspace(0, 1.25, size).unsqueeze(0),
torch.linspace(0, 1.5, size).unsqueeze(0),
torch.linspace(0, 1, size).unsqueeze(0),
torch.linspace(0, 0.5, size).unsqueeze(0),
torch.linspace(0, 0.25, size).unsqueeze(0),
torch.linspace(0, 1.25, size).unsqueeze(0),
torch.linspace(0, 1.25, size).unsqueeze(0),
torch.linspace(0, 1.5, size).unsqueeze(0),
torch.linspace(0, 1, size).unsqueeze(0),
],
0,
).unsqueeze(-1)
covar_matrix = RBFKernel()(train_x, train_x).evaluate().view(
2, 2, 3, size, size)
piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
covar_approx = piv_chol.transpose(-1, -2).matmul(piv_chol)
self.assertTrue(approx_equal(covar_approx, covar_matrix, 2e-4))
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])
)
def test_solve(self):
size = 100
train_x = torch.linspace(0, 1, size)
covar_matrix = RBFKernel()(train_x, train_x).evaluate()
piv_chol = pivoted_cholesky.pivoted_cholesky(covar_matrix, 10)
woodbury_factor, inv_scale, logdet = woodbury.woodbury_factor(piv_chol, piv_chol, torch.ones(100), logdet=True)
self.assertTrue(approx_equal(logdet, (piv_chol @ piv_chol.transpose(-1, -2) + torch.eye(100)).logdet(), 2e-4))
rhs_vector = torch.randn(100, 50)
shifted_covar_matrix = covar_matrix + torch.eye(size)
real_solve = shifted_covar_matrix.inverse().matmul(rhs_vector)
scaled_inv_diag = (inv_scale / torch.ones(100)).unsqueeze(-1)
approx_solve = woodbury.woodbury_solve(
rhs_vector, piv_chol * scaled_inv_diag, woodbury_factor, scaled_inv_diag, inv_scale
)
self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
def test_inv_quad_only_vector(self):
# Forward pass
res = NonLazyTensor(self.mat_var).inv_quad(self.vec_var)
actual = self.mat_var_clone.inverse().matmul(self.vec_var_clone).mul(
self.vec_var_clone).sum()
self.assertAlmostEqual(res.item(), actual.item(), places=1)
# Backward
actual.backward()
res.backward()
self.assertTrue(
approx_equal(self.mat_var_clone.grad,
self.mat_var.grad,
epsilon=1e-1))
self.assertTrue(
approx_equal(self.vec_var_clone.grad, self.vec_var.grad))
def test_permute(self):
lazy_tensor = self.create_lazy_tensor()
if lazy_tensor.dim() >= 4:
evaluated = self.evaluate_lazy_tensor(lazy_tensor)
dims = torch.randperm(lazy_tensor.dim() - 2).tolist()
res = lazy_tensor.permute(*dims, -2, -1).evaluate()
actual = evaluated.permute(*dims, -2, -1)
self.assertTrue(approx_equal(res, actual))
def test_inv_matmul_multiple_vecs(self):
# Forward
res = NonLazyTensor(self.mats_var).inv_matmul(self.vecs_var)
actual = torch.cat([
self.mats_var_copy[0].inverse().unsqueeze(0),
self.mats_var_copy[1].inverse().unsqueeze(0)
]).matmul(self.vecs_var_copy)
self.assertTrue(approx_equal(res, actual))
# Backward
grad_output = torch.randn(2, 3, 4)
res.backward(gradient=grad_output)
actual.backward(gradient=grad_output)
self.assertTrue(
approx_equal(self.mats_var_copy.grad, self.mats_var.grad))
self.assertTrue(
approx_equal(self.vecs_var_copy.grad, self.vecs_var.grad))
def test_add_diag(self):
diag = torch.tensor(1.5)
res = ZeroLazyTensor(5, 5).add_diag(diag).evaluate()
actual = torch.eye(5).mul(1.5)
self.assertTrue(approx_equal(res, actual))
diag = torch.tensor([1.5])
res = ZeroLazyTensor(5, 5).add_diag(diag).evaluate()
actual = torch.eye(5).mul(1.5)
self.assertTrue(approx_equal(res, actual))
diag = torch.tensor([1.5, 1.3, 1.2, 1.1, 2.])
res = ZeroLazyTensor(5, 5).add_diag(diag).evaluate()
actual = diag.diag()
self.assertTrue(approx_equal(res, actual))
diag = torch.tensor(1.5)
res = ZeroLazyTensor(2, 5, 5).add_diag(diag).evaluate()
actual = torch.eye(5).unsqueeze(0).repeat(2, 1, 1).mul(1.5)
self.assertTrue(approx_equal(res, actual))
diag = torch.tensor([1.5])
res = ZeroLazyTensor(2, 5, 5).add_diag(diag).evaluate()
actual = torch.eye(5).unsqueeze(0).repeat(2, 1, 1).mul(1.5)
self.assertTrue(approx_equal(res, actual))
diag = torch.tensor([1.5, 1.3, 1.2, 1.1, 2.])
res = ZeroLazyTensor(2, 5, 5).add_diag(diag).evaluate()
actual = diag.diag().unsqueeze(0).repeat(2, 1, 1)
self.assertTrue(approx_equal(res, actual))
diag = torch.tensor([[1.5, 1.3, 1.2, 1.1, 2.], [0, 1, 2, 1, 1]])
res = ZeroLazyTensor(2, 5, 5).add_diag(diag).evaluate()
actual = torch.cat([diag[0].diag().unsqueeze(0), diag[1].diag().unsqueeze(0)])
self.assertTrue(approx_equal(res, actual))
def test_multivariate_normal_batch_non_lazy(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
mean = torch.tensor([0, 1, 2], dtype=torch.float, device=device)
covmat = torch.diag(torch.tensor([1, 0.75, 1.5], device=device))
mvn = MultivariateNormal(mean=mean.repeat(2, 1), covariance_matrix=covmat.repeat(2, 1, 1), validate_args=True)
self.assertTrue(torch.is_tensor(mvn.covariance_matrix))
self.assertIsInstance(mvn.lazy_covariance_matrix, LazyTensor)
self.assertTrue(approx_equal(mvn.variance, covmat.diag().repeat(2, 1)))
self.assertTrue(approx_equal(mvn.scale_tril, torch.diag(covmat.diag().sqrt()).repeat(2, 1, 1)))
mvn_plus1 = mvn + 1
self.assertTrue(torch.equal(mvn_plus1.mean, mvn.mean + 1))
self.assertTrue(torch.equal(mvn_plus1.covariance_matrix, mvn.covariance_matrix))
mvn_times2 = mvn * 2
self.assertTrue(torch.equal(mvn_times2.mean, mvn.mean * 2))
self.assertTrue(torch.equal(mvn_times2.covariance_matrix, mvn.covariance_matrix * 4))
mvn_divby2 = mvn / 2
self.assertTrue(torch.equal(mvn_divby2.mean, mvn.mean / 2))
self.assertTrue(torch.equal(mvn_divby2.covariance_matrix, mvn.covariance_matrix / 4))
self.assertTrue(approx_equal(mvn.entropy(), 4.3157 * torch.ones(2, device=device)))
logprob = mvn.log_prob(torch.zeros(2, 3, device=device))
logprob_expected = -4.8157 * torch.ones(2, device=device)
self.assertTrue(approx_equal(logprob, logprob_expected))
logprob = mvn.log_prob(torch.zeros(2, 2, 3, device=device))
logprob_expected = -4.8157 * torch.ones(2, 2, device=device)
self.assertTrue(approx_equal(logprob, logprob_expected))
conf_lower, conf_upper = mvn.confidence_region()
self.assertTrue(approx_equal(conf_lower, mvn.mean - 2 * mvn.stddev))
self.assertTrue(approx_equal(conf_upper, mvn.mean + 2 * mvn.stddev))
self.assertTrue(mvn.sample().shape == torch.Size([2, 3]))
self.assertTrue(mvn.sample(torch.Size([2])).shape == torch.Size([2, 2, 3]))
self.assertTrue(mvn.sample(torch.Size([2, 4])).shape == torch.Size([2, 4, 2, 3]))
def test_add_diag(self):
lazy_tensor = self.create_lazy_tensor()
evaluated = self.evaluate_lazy_tensor(lazy_tensor)
other_diag = torch.tensor(1.5)
res = lazy_tensor.add_diag(other_diag).evaluate()
actual = evaluated + torch.eye(evaluated.size(-1)).mul(1.5)
self.assertTrue(approx_equal(res, actual))
other_diag = torch.tensor([1.5])
res = lazy_tensor.add_diag(other_diag).evaluate()
actual = evaluated + torch.eye(evaluated.size(-1)).mul(1.5)
self.assertTrue(approx_equal(res, actual))
other_diag = torch.randn(lazy_tensor.size(-1)).pow(2)
res = lazy_tensor.add_diag(other_diag).evaluate()
actual = evaluated + other_diag.diag()
self.assertTrue(approx_equal(res, actual))
def test_lkj_covariance_prior_batch_log_prob(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
v = torch.ones(2, 1, device=device)
sd_prior = SmoothedBoxPrior(exp(-1) * v, exp(1) * v)
prior = LKJCovariancePrior(2, torch.tensor([0.5, 1.5], device=device),
sd_prior)
S = torch.eye(2, device=device)
self.assertTrue(
approx_equal(prior.log_prob(S),
torch.tensor([-3.59981, -2.21351], device=S.device)))
S = torch.stack(
[S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
self.assertTrue(
approx_equal(prior.log_prob(S),
torch.tensor([-3.59981, -2.35735], device=S.device)))
with self.assertRaises(ValueError):
prior.log_prob(torch.eye(3, device=device))
def test_log_det_only(self):
# Forward pass
with gpytorch.settings.num_trace_samples(1000):
res = NonLazyTensor(self.mats_var).log_det()
actual = torch.cat([
self.mats_var_clone[0].det().log().unsqueeze(0),
self.mats_var_clone[1].det().log().unsqueeze(0)
])
self.assertTrue(approx_equal(res, actual, epsilon=1e-1))
# Backward
grad_output = torch.tensor([3, 4], dtype=torch.float)
actual.backward(gradient=grad_output)
res.backward(gradient=grad_output)
self.assertTrue(
approx_equal(self.mats_var_clone.grad,
self.mats_var.grad,
epsilon=1e-1))
def test_lkj_covariance_prior_log_prob_hetsd(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
a = torch.tensor([exp(-1), exp(-2)], device=device)
b = torch.tensor([exp(1), exp(2)], device=device)
sd_prior = SmoothedBoxPrior(a, b, log_transform=True)
prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior)
self.assertFalse(prior.log_transform)
S = torch.eye(2, device=device)
self.assertAlmostEqual(prior.log_prob(S).item(), -4.71958, places=4)
S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
self.assertTrue(approx_equal(prior.log_prob(S), torch.tensor([-4.71958, -4.57574], device=S.device)))
with self.assertRaises(ValueError):
prior.log_prob(torch.eye(3, device=device))
# For eta=1.0 log_prob is flat over all covariance matrices
prior = LKJCovariancePrior(2, torch.tensor(1.0, device=device), sd_prior)
marginal_sd = torch.diagonal(S, dim1=-2, dim2=-1).sqrt()
log_prob_expected = prior.correlation_prior.C + prior.sd_prior.log_prob(marginal_sd)
self.assertTrue(approx_equal(prior.log_prob(S), log_prob_expected))
def test_solve_qr_constant_noise(self, dtype=torch.float64, tol=1e-8):
size = 50
X = torch.rand((size, 2)).to(dtype=dtype)
y = torch.sin(torch.sum(X, 1)).unsqueeze(-1).to(dtype=dtype)
noise = 1e-2 * torch.ones(size, dtype=dtype)
lazy_tsr = RBFKernel().to(dtype=dtype)(X).evaluate_kernel().add_diag(noise)
precondition_qr, _, logdet_qr = lazy_tsr._preconditioner()
F = lazy_tsr._piv_chol_self
M = noise.diag() + F.matmul(F.t())
x_exact = torch.solve(y, M)[0]
x_qr = precondition_qr(y)
self.assertTrue(approx_equal(x_exact, x_qr, tol))
logdet = 2 * torch.cholesky(M).diag().log().sum(-1)
self.assertTrue(approx_equal(logdet, logdet_qr, tol))
def test_root_inv_decomposition(self):
# Forward
probe_vectors = torch.randn(2, 3, 4, 5)
test_vectors = torch.randn(2, 3, 4, 5)
root = NonLazyTensor(self.mat).root_inv_decomposition(
probe_vectors, test_vectors).root.evaluate()
res = root.matmul(root.transpose(-1, -2))
flattened_mats = self.mat_clone.view(-1, *self.mat_clone.shape[-2:])
actual = torch.cat([
mat.inverse().unsqueeze(0) for mat in flattened_mats
]).view_as(self.mat_clone)
self.assertTrue(approx_equal(res, actual))
# Backward
sum([mat.trace()
for mat in res.view(-1, *self.mat.shape[-2:])]).backward()
sum([mat.trace()
for mat in actual.view(-1, *self.mat.shape[-2:])]).backward()
self.assertTrue(approx_equal(self.mat.grad, self.mat_clone.grad))
def test_inv_quad_only_many_vectors(self):
# Forward pass
res = NonLazyTensor(self.mats_var).inv_quad(self.vecs_var).sum()
actual = (torch.cat([
self.mats_var_clone[0].inverse().unsqueeze(0),
self.mats_var_clone[1].inverse().unsqueeze(0)
]).matmul(self.vecs_var_clone).mul(
self.vecs_var_clone).sum(2).sum(1)).sum()
self.assertTrue(approx_equal(res, actual, epsilon=1e-1))
# Backward
actual.backward()
res.backward()
self.assertTrue(
approx_equal(self.mats_var_clone.grad,
self.mats_var.grad,
epsilon=1e-1))
self.assertTrue(
approx_equal(self.vecs_var_clone.grad, self.vecs_var.grad))
def test_multitask_multivariate_normal(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
mean = torch.tensor([[0, 1], [2, 3]], dtype=torch.float, device=device)
variance = 1 + torch.arange(4, dtype=torch.float, device=device)
covmat = torch.diag(variance)
mtmvn = MultitaskMultivariateNormal(mean=mean,
covariance_matrix=covmat)
self.assertTrue(torch.equal(mtmvn.mean, mean))
self.assertTrue(approx_equal(mtmvn.variance, variance.view(2, 2)))
self.assertTrue(torch.equal(mtmvn.scale_tril, covmat.sqrt()))
self.assertTrue(mtmvn.event_shape == torch.Size([2, 2]))
mvn_plus1 = mtmvn + 1
self.assertTrue(torch.equal(mvn_plus1.mean, mtmvn.mean + 1))
self.assertTrue(
torch.equal(mvn_plus1.covariance_matrix, mtmvn.covariance_matrix))
mvn_times2 = mtmvn * 2
self.assertTrue(torch.equal(mvn_times2.mean, mtmvn.mean * 2))
self.assertTrue(
torch.equal(mvn_times2.covariance_matrix,
mtmvn.covariance_matrix * 4))
mvn_divby2 = mtmvn / 2
self.assertTrue(torch.equal(mvn_divby2.mean, mtmvn.mean / 2))
self.assertTrue(
torch.equal(mvn_divby2.covariance_matrix,
mtmvn.covariance_matrix / 4))
self.assertAlmostEqual(mtmvn.entropy().item(), 7.2648, places=4)
self.assertAlmostEqual(mtmvn.log_prob(torch.zeros(
2, 2, device=device)).item(),
-7.3064,
places=4)
logprob = mtmvn.log_prob(torch.zeros(3, 2, 2, device=device))
logprob_expected = -7.3064 * torch.ones(3, device=device)
self.assertTrue(approx_equal(logprob, logprob_expected))
conf_lower, conf_upper = mtmvn.confidence_region()
self.assertTrue(approx_equal(conf_lower,
mtmvn.mean - 2 * mtmvn.stddev))
self.assertTrue(approx_equal(conf_upper,
mtmvn.mean + 2 * mtmvn.stddev))
self.assertTrue(mtmvn.sample().shape == torch.Size([2, 2]))
self.assertTrue(
mtmvn.sample(torch.Size([3])).shape == torch.Size([3, 2, 2]))
self.assertTrue(
mtmvn.sample(torch.Size([3, 4])).shape == torch.Size([3, 4, 2, 2]))
def test_multivariate_normal_batch_lazy(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
mean = torch.tensor([0, 1, 2], dtype=torch.float, device=device)
covmat = torch.diag(torch.tensor([1, 0.75, 1.5], device=device))
mvn = MultivariateNormal(
mean=mean.repeat(2, 1),
covariance_matrix=NonLazyTensor(covmat).repeat(2, 1, 1))
self.assertTrue(torch.is_tensor(mvn.covariance_matrix))
self.assertIsInstance(mvn.lazy_covariance_matrix, LazyTensor)
self.assertTrue(approx_equal(mvn.variance, covmat.diag().repeat(2, 1)))
mvn_plus1 = mvn + 1
self.assertTrue(torch.equal(mvn_plus1.mean, mvn.mean + 1))
self.assertTrue(
torch.equal(mvn_plus1.covariance_matrix, mvn.covariance_matrix))
mvn_times2 = mvn * 2
self.assertTrue(torch.equal(mvn_times2.mean, mvn.mean * 2))
self.assertTrue(
torch.equal(mvn_times2.covariance_matrix,
mvn.covariance_matrix * 4))
mvn_divby2 = mvn / 2
self.assertTrue(torch.equal(mvn_divby2.mean, mvn.mean / 2))
self.assertTrue(
torch.equal(mvn_divby2.covariance_matrix,
mvn.covariance_matrix / 4))
# TODO: Add tests for entropy, log_prob, etc. - this an issue b/c it
# uses using root_decomposition which is not very reliable
# self.assertTrue(approx_equal(mvn.entropy(), 4.3157 * torch.ones(2)))
# self.assertTrue(
# approx_equal(mvn.log_prob(torch.zeros(2, 3)), -4.8157 * torch.ones(2))
# )
# self.assertTrue(
# approx_equal(mvn.log_prob(torch.zeros(2, 2, 3)), -4.8157 * torch.ones(2, 2))
# )
conf_lower, conf_upper = mvn.confidence_region()
self.assertTrue(approx_equal(conf_lower, mvn.mean - 2 * mvn.stddev))
self.assertTrue(approx_equal(conf_upper, mvn.mean + 2 * mvn.stddev))
self.assertTrue(mvn.sample().shape == torch.Size([2, 3]))
self.assertTrue(
mvn.sample(torch.Size([2])).shape == torch.Size([2, 2, 3]))
self.assertTrue(
mvn.sample(torch.Size([2, 4])).shape == torch.Size([2, 4, 2, 3]))
def test_lanczos(self):
size = 100
matrix = torch.randn(size, size)
matrix = matrix.matmul(matrix.transpose(-1, -2))
matrix.div_(matrix.norm())
matrix.add_(torch.ones(matrix.size(-1)).mul(1e-6).diag())
q_mat, t_mat = lanczos_tridiag(
matrix.matmul, max_iter=size, dtype=matrix.dtype, device=matrix.device, matrix_shape=matrix.shape
)
approx = q_mat.matmul(t_mat).matmul(q_mat.transpose(-1, -2))
self.assertTrue(approx_equal(approx, matrix))
def test_lkj_cholesky_factor_prior_batch_log_prob(self, cuda=False):
device = torch.device("cuda") if cuda else torch.device("cpu")
prior = LKJCholeskyFactorPrior(2,
torch.tensor([0.5, 1.5], device=device))
self.assertFalse(prior.log_transform)
S = torch.eye(2, device=device)
S_chol = torch.potrf(S, upper=False)
self.assertTrue(
approx_equal(
prior.log_prob(S_chol),
torch.tensor([-1.86942, -0.483129], device=S_chol.device)))
S = torch.stack(
[S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
S_chol = torch.stack([torch.potrf(Si, upper=False) for Si in S])
self.assertTrue(
approx_equal(
prior.log_prob(S_chol),
torch.tensor([-1.86942, -0.62697], device=S_chol.device)))
with self.assertRaises(ValueError):
prior.log_prob(torch.eye(3, device=device))
def test_solve(self):
size = 100
train_x = torch.linspace(0, 1, size)
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.ones(100))
rhs_vector = torch.randn(100, 50)
shifted_covar_matrix = covar_matrix + torch.eye(size)
real_solve = shifted_covar_matrix.inverse().matmul(rhs_vector)
approx_solve = pivoted_cholesky.woodbury_solve(rhs_vector, piv_chol, woodbury_factor, torch.ones(100))
self.assertTrue(approx_equal(approx_solve, real_solve, 2e-4))
