Repository: cornellius-gp/gpytorch
Pull request #305 (2018-09-30 18:44:09)

Make LKJCovariancePrior work properly in batch mode
Makes sure LKJCovariancePrior works properly in batch mode.
Also adds the ability to have heterogeneous priors over the different
marginal variances (e.g. can use a sd_prior with event_shape that is
either 1 or n if the covariance matrices are n x n)


Patch:
diff --git a/gpytorch/priors/lkj_prior.py b/gpytorch/priors/lkj_prior.py
index 4c659034..041a39ca 100644
--- a/gpytorch/priors/lkj_prior.py
+++ b/gpytorch/priors/lkj_prior.py
@@ -58,12 +58,7 @@ class LKJPrior(Prior):
             raise ValueError("Input is not a valid correlation matrix")
         # TODO: Replace this loop with batch Cholesky decomposition when available
         # https://github.com/pytorch/pytorch/pull/11796
-        log_diag_sum = torch.stack(
-            [
-                p.potrf().diag().log().sum()
-                for p in parameter.view(-1, *parameter.shape[-2:])
-            ]
-        )
+        log_diag_sum = torch.stack([p.potrf().diag().log().sum() for p in parameter.view(-1, *parameter.shape[-2:])])
         return self.C + (self.eta - 1) * 2 * log_diag_sum
 
 
@@ -85,6 +80,8 @@ class LKJCholeskyFactorPrior(LKJPrior):
     def log_prob(self, parameter):
         if self.log_transform:
             raise RuntimeError("log-transform not supported for LKJCholeskyFactorPrior")
+        if any(s != self.n for s in parameter.shape[-2:]):
+            raise ValueError("Cholesky factor is not of size n={}".format(self.n.item()))
         if not _is_valid_correlation_matrix_cholesky_factor(parameter):
             raise ValueError("Input is not a Cholesky factor of a valid correlation matrix")
         log_diag_sum = torch.diagonal(parameter, dim1=-2, dim2=-1).log().sum(-1)
@@ -106,16 +103,16 @@ class LKJCovariancePrior(LKJPrior):
     def __init__(self, n, eta, sd_prior, validate_args=False):
         if not isinstance(sd_prior, Prior):
             raise ValueError("sd_prior must be an instance of Prior")
-        if sd_prior.event_shape != torch.Size([1]):
-            raise ValueError("Currently only sd_prior w/ event_shape=1 is supported")
-        if sd_prior.batch_shape != torch.Size():
-            raise ValueError("sd_prior must have same batch_shape as eta")
+        if not isinstance(n, int):
+            raise ValueError("n must be an integer")
+        if sd_prior.event_shape not in {torch.Size([1]), torch.Size([n])}:
+            raise ValueError("sd_prior must have event_shape 1 or n")
         correlation_prior = LKJPrior(n=n, eta=eta, validate_args=validate_args)
+        if sd_prior.batch_shape != correlation_prior.batch_shape:
+            raise ValueError("sd_prior must have same batch_shape as eta")
         TModule.__init__(self)
         super(LKJPrior, self).__init__(
-            correlation_prior.batch_shape,
-            correlation_prior.event_shape,
-            validate_args=validate_args,
+            correlation_prior.batch_shape, correlation_prior.event_shape, validate_args=False
         )
         self.correlation_prior = correlation_prior
         self.sd_prior = sd_prior
@@ -125,23 +122,19 @@ class LKJCovariancePrior(LKJPrior):
         if self.log_transform:
             raise RuntimeError("log-transform not supported for LKJCovariancePrior")
         marginal_var = torch.diagonal(parameter, dim1=-2, dim2=-1)
-        if torch.all(marginal_var >= 0):
-            marginal_sd = marginal_var.sqrt()
-        else:
+        if not torch.all(marginal_var >= 0):
             raise ValueError("Variance(s) cannot be negative")
+        marginal_sd = marginal_var.sqrt()
         sd_diag_mat = _batch_form_diag(1 / marginal_sd)
-        mm = torch.bmm if parameter.dim() > 2 else torch.mm
-        correlations = mm(torch.mm(sd_diag_mat, parameter), sd_diag_mat)
-        log_prob = self.correlation_prior.log_prob(correlations)
-        # Add log likelihoods of each of the n marginal standard deviations
-        for i in range(self.correlation_prior.n):
-            log_prob = log_prob + self.sd_prior.log_prob(marginal_sd[i])
-        return log_prob
+        correlations = torch.matmul(torch.matmul(sd_diag_mat, parameter), sd_diag_mat)
+        log_prob_corr = self.correlation_prior.log_prob(correlations)
+        log_prob_sd = self.sd_prior.log_prob(marginal_sd)
+        return log_prob_corr + log_prob_sd
 
 
 def _batch_form_diag(tsr):
     """Form diagonal matrices in batch mode."""
-    eye = torch.eye(tsr.shape[-1])
+    eye = torch.eye(tsr.shape[-1], dtype=tsr.dtype, device=tsr.device)
     M = tsr.unsqueeze(-1).expand(tsr.shape + tsr.shape[-1:])
     return eye * M
 
@@ -163,10 +156,7 @@ def _is_valid_correlation_matrix(Sigma, tol=1e-6):
 
     """
     pdef = torch.all(constraints.positive_definite.check(Sigma))
-    return pdef and all(
-        torch.all(torch.abs(S.diag() - 1) < tol)
-        for S in Sigma.view(-1, *Sigma.shape[-2:])
-    )
+    return pdef and all(torch.all(torch.abs(S.diag() - 1) < tol) for S in Sigma.view(-1, *Sigma.shape[-2:]))
 
 
 def _is_valid_correlation_matrix_cholesky_factor(L, tol=1e-6):
diff --git a/gpytorch/priors/smoothed_box_prior.py b/gpytorch/priors/smoothed_box_prior.py
index 049b278a..cccef639 100644
--- a/gpytorch/priors/smoothed_box_prior.py
+++ b/gpytorch/priors/smoothed_box_prior.py
@@ -24,7 +24,7 @@ class SmoothedBoxPrior(Prior):
 
     """
 
-    arg_constraints = {"sigma": constraints.positive}
+    arg_constraints = {"sigma": constraints.positive, "a": constraints.real, "b": constraints.real}
     support = constraints.real
     _validate_args = True
 
@@ -37,16 +37,16 @@ class SmoothedBoxPrior(Prior):
             raise ValueError("must have that a < b (element-wise)")
         # TODO: Proper argument validation including broadcasting
         batch_shape, event_shape = _a.shape[:-1], _a.shape[-1:]
+        # need to assign values before registering as buffers to make argument validation work
+        self.a, self.b, self.sigma = _a, _b, _sigma
+        super(SmoothedBoxPrior, self).__init__(batch_shape, event_shape, validate_args=validate_args)
+        # now need to delete to be able to register buffer
+        del self.a, self.b, self.sigma
         self.register_buffer("a", _a)
         self.register_buffer("b", _b)
         self.register_buffer("sigma", _sigma)
-        self.tails = NormalPrior(
-            torch.zeros_like(_a), _sigma, validate_args=validate_args
-        )
+        self.tails = NormalPrior(torch.zeros_like(_a), _sigma, validate_args=validate_args)
         self._log_transform = log_transform
-        super(SmoothedBoxPrior, self).__init__(
-            batch_shape, event_shape, validate_args=validate_args
-        )
 
     @property
     def _c(self):
@@ -59,9 +59,7 @@ class SmoothedBoxPrior(Prior):
     @property
     def _M(self):
         # normalization factor to make this a probability distribution
-        return torch.log(
-            1 + (self.b - self.a) / (math.sqrt(2 * math.pi) * self.sigma)
-        )
+        return torch.log(1 + (self.b - self.a) / (math.sqrt(2 * math.pi) * self.sigma))
 
     def log_prob(self, parameter):
         return self._log_prob(parameter.exp() if self.log_transform else parameter)


Test patch:
diff --git a/test/priors/test_lkj_prior.py b/test/priors/test_lkj_prior.py
index 9edcaedd..2df3d857 100644
--- a/test/priors/test_lkj_prior.py
+++ b/test/priors/test_lkj_prior.py
@@ -1,9 +1,10 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import unittest
+from math import exp
 
 import torch
-from gpytorch.priors import LKJPrior
+from gpytorch.priors import LKJCholeskyFactorPrior, LKJCovariancePrior, LKJPrior, SmoothedBoxPrior
 from gpytorch.utils import approx_equal
 
 
@@ -58,5 +59,143 @@ class TestLKJPrior(unittest.TestCase):
             return self.test_lkj_prior_batch_log_prob(cuda=True)
 
 
+class TestLKJCholeskyFactorPrior(unittest.TestCase):
+    def test_lkj_cholesky_factor_prior_to_gpu(self):
+        if torch.cuda.is_available():
+            prior = LKJCholeskyFactorPrior(2, 1.0).cuda()
+            self.assertEqual(prior.eta.device.type, "cuda")
+            self.assertEqual(prior.C.device.type, "cuda")
+
+    def test_lkj_cholesky_factor_prior_validate_args(self):
+        LKJCholeskyFactorPrior(2, 1.0, validate_args=True)
+        with self.assertRaises(ValueError):
+            LKJCholeskyFactorPrior(1.5, 1.0, validate_args=True)
+        with self.assertRaises(ValueError):
+            LKJCholeskyFactorPrior(2, -1.0, validate_args=True)
+
+    def test_lkj_cholesky_factor_prior_log_prob(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        prior = LKJCholeskyFactorPrior(2, torch.tensor(0.5, device=device))
+        self.assertFalse(prior.log_transform)
+        S = torch.eye(2, device=device)
+        S_chol = torch.potrf(S, upper=False)
+        self.assertAlmostEqual(prior.log_prob(S_chol).item(), -1.86942, places=4)
+        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S_chol.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, -1.72558], device=S_chol.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 = LKJCholeskyFactorPrior(2, torch.tensor(1.0, device=device))
+        self.assertTrue(torch.all(prior.log_prob(S_chol) == prior.C))
+
+    def test_lkj_cholesky_factor_prior_log_prob_cuda(self):
+        if torch.cuda.is_available():
+            return self.test_lkj_cholesky_factor_prior_log_prob(cuda=True)
+
+    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_lkj_cholesky_factor_prior_batch_log_prob_cuda(self):
+        if torch.cuda.is_available():
+            return self.test_lkj_cholesky_factor_prior_batch_log_prob(cuda=True)
+
+
+class TestLKJCovariancePrior(unittest.TestCase):
+    def test_lkj_covariance_prior_to_gpu(self):
+        if torch.cuda.is_available():
+            sd_prior = SmoothedBoxPrior(exp(-1), exp(1), log_transform=True)
+            prior = LKJCovariancePrior(2, 1.0, sd_prior).cuda()
+            self.assertEqual(prior.correlation_prior.eta.device.type, "cuda")
+            self.assertEqual(prior.correlation_prior.C.device.type, "cuda")
+            self.assertEqual(prior.sd_prior.a.device.type, "cuda")
+
+    def test_lkj_covariance_prior_validate_args(self):
+        sd_prior = SmoothedBoxPrior(exp(-1), exp(1), log_transform=True, validate_args=True)
+        LKJCovariancePrior(2, 1.0, sd_prior)
+        with self.assertRaises(ValueError):
+            LKJCovariancePrior(1.5, 1.0, sd_prior, validate_args=True)
+        with self.assertRaises(ValueError):
+            LKJCovariancePrior(2, -1.0, sd_prior, validate_args=True)
+
+    def test_lkj_covariance_prior_log_prob(self, cuda=False):
+        device = torch.device("cuda") if cuda else torch.device("cpu")
+        sd_prior = SmoothedBoxPrior(exp(-1), exp(1), log_transform=True)
+        if cuda:
+            sd_prior = sd_prior.cuda()
+        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(), -3.59981, 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([-3.59981, -3.45597], 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_lkj_covariance_prior_log_prob_cuda(self):
+        if torch.cuda.is_available():
+            return self.test_lkj_covariance_prior_log_prob(cuda=True)
+
+    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_lkj_covariance_prior_log_prob_hetsd_cuda(self):
+        if torch.cuda.is_available():
+            return self.test_lkj_covariance_prior_log_prob_hetsd(cuda=True)
+
+    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, log_transform=True)
+        prior = LKJCovariancePrior(2, torch.tensor([0.5, 1.5], device=device), sd_prior)
+
+        self.assertFalse(prior.log_transform)
+        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_lkj_covariance_prior_batch_log_prob_cuda(self):
+        if torch.cuda.is_available():
+            return self.test_lkj_covariance_prior_batch_log_prob(cuda=True)
+
+
 if __name__ == "__main__":
     unittest.main()