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))