Repository: cornellius-gp/gpytorch

Issue:
Trying to sample from the LKJCovariancePrior
Trying to declare and LKJ prior and sample from it however it doesn't like the event_shape of the sd_prior:

```
import torch
from gpytorch.priors import LKJCovariancePrior, GammaPrior

sd_prior = GammaPrior(torch.tensor([1.0]),torch.tensor([1.0]), validate_args=True)
cov_prior = LKJCovariancePrior(n=10, eta=1, sd_prior=sd_prior)
```
```
Traceback (most recent call last):

  File "<ipython-input-40-1ca20a6657f1>", line 5, in <module>
    cov_prior = LKJCovariancePrior(n=10, eta=1, sd_prior=sd_prior)

  File "/home/vr308/anaconda3/lib/python3.8/site-packages/gpytorch/priors/lkj_prior.py", line 111, in __init__
    raise ValueError("sd_prior must have event_shape 1 or n")

ValueError: sd_prior must have event_shape 1 or n

```
event_shape is not a settable attribute and by default it is empty torch.Size([]) for a prior over scalar - is this a bug?

Thanks.


Discussion:
Hi, I was dealing with this the other day, and don't think it's actually a bug in gpytorch but rather a bug in `torch.distributions.ExponentialFamily` (or some strange feature), as there's no way to set event shape in those distributions.

This works though because you can set `_event_shape`:
```python
sd_prior = GammaPrior(torch.tensor([1.0]),torch.tensor([1.0]), validate_args=True)
sd_prior._event_shape = torch.Size((1,))
cov_prior = LKJCovariancePrior(n=10, eta=torch.ones(1), sd_prior=sd_prior)
```
thanks, but, how do we get the sampling to work?

`cov_prior.sample()
`

throws a notImplemented error.

We need to write a `sample()` implementation for the LKJPrior. Open to a PR! ;)
I can try to setup a PR soon, finally managed to track down some code for sampling from the LKJ priors: 

- Python (numpy): https://github.com/pymc-devs/pymc3/blob/7e464e59bcb0adb28df94f379b3e8d4af12bd4d1/pymc3/distributions/multivariate.py#L1107 
- R (this is much more readable): https://github.com/rmcelreath/rethinking/blob/2acf2fd7b01718cf66a8352c52d001886c7d3c4c/R/distributions.r#L214

Resolving patch:
diff --git a/gpytorch/priors/lkj_prior.py b/gpytorch/priors/lkj_prior.py
--- a/gpytorch/priors/lkj_prior.py
+++ b/gpytorch/priors/lkj_prior.py
@@ -1,10 +1,7 @@
 #!/usr/bin/env python3
 
-import math
-from numbers import Number
-
 import torch
-from torch.distributions import constraints
+from torch.distributions import LKJCholesky, constraints
 from torch.nn import Module as TModule
 
 from .. import settings
@@ -12,8 +9,9 @@
 from .prior import Prior
 
 
-class LKJPrior(Prior):
-    r"""LKJ prior over n x n (positive definite) correlation matrices
+class LKJCholeskyFactorPrior(Prior, LKJCholesky):
+    r"""LKJ prior over n x n (positive definite) Cholesky-decomposed
+    correlation matrices
 
     .. math:
 
@@ -21,50 +19,27 @@ class LKJPrior(Prior):
             pdf(\Sigma) ~ |\Sigma| ^ (\eta  - 1)
         \end{equation*}
 
-    where :math:`\eta > 0` is a shape parameter.
+    where :math:`\eta > 0` is a shape parameter and n is the dimension of the
+    correlation matrix.
 
-    Reference: Bayesian Data Analysis, 3rd ed., Gelman et al., p. 576
-    """
+    LKJCholeskyFactorPrior is different from LKJPrior in that it accepts the
+    Cholesky factor of the correlation matrix to compute probabilities.
 
-    arg_constraints = {"n": constraints.positive_integer, "eta": constraints.positive}
-    # TODO: move correlation matrix validation upstream into pytorch
-    support = constraints.positive_definite
-    _validate_args = True
+    This distribution assumes that the cholesky factor is lower-triangular.
+    """
 
-    def __init__(self, n, eta, validate_args=False):
+    def __init__(self, n, eta, validate_args=False, transform=None):
         TModule.__init__(self)
-        if not isinstance(n, int) or n < 1:
-            raise ValueError("n must be a positive integer")
-        if isinstance(eta, Number):
-            eta = torch.tensor(float(eta))
-        self.n = torch.tensor(n, dtype=torch.long, device=eta.device)
-        batch_shape = eta.shape
-        event_shape = torch.Size([n, n])
-        # Normalization constant(s)
-        i = torch.arange(n, dtype=eta.dtype, device=eta.device)
-        C = (((2 * eta.view(-1, 1) - 2 + i) * i).sum(1) * math.log(2)).view_as(eta)
-        C += n * torch.sum(2 * torch.lgamma(i / 2 + 1) - torch.lgamma(i + 2))
-        # need to assign values before registering as buffers to make argument validation work
-        self.eta = eta
-        self.C = C
-        super(LKJPrior, self).__init__(batch_shape, event_shape, validate_args=validate_args)
-        # now need to delete to be able to register buffer
-        del self.eta, self.C
-        self.register_buffer("eta", eta)
-        self.register_buffer("C", C)
+        LKJCholesky.__init__(self, dim=n, concentration=eta, validate_args=validate_args)
+        self.n = self.dim
+        self.eta = self.concentration
+        self._transform = transform
 
-    def log_prob(self, X):
-        if any(s != self.n for s in X.shape[-2:]):
-            raise ValueError("Correlation matrix is not of size n={}".format(self.n.item()))
-        if not _is_valid_correlation_matrix(X):
-            raise ValueError("Input is not a valid correlation matrix")
-        log_diag_sum = psd_safe_cholesky(X, upper=True).diagonal(dim1=-2, dim2=-1).log().sum(-1)
-        return self.C + (self.eta - 1) * 2 * log_diag_sum
+    # TODO: change from sample to rsample if pytorch #69281 goes in providing rsample method
 
 
-class LKJCholeskyFactorPrior(LKJPrior):
-    r"""LKJ prior over n x n (positive definite) Cholesky-decomposed
-    correlation matrices
+class LKJPrior(LKJCholeskyFactorPrior):
+    r"""LKJ prior over n x n (positive definite) correlation matrices
 
     .. math:
 
@@ -72,22 +47,23 @@ class LKJCholeskyFactorPrior(LKJPrior):
             pdf(\Sigma) ~ |\Sigma| ^ (\eta  - 1)
         \end{equation*}
 
-    where :math:`\eta > 0` is a shape parameter and n is the dimension of the
-    correlation matrix.
+    where :math:`\eta > 0` is a shape parameter.
 
-    LKJCholeskyFactorPrior is different from LKJPrior in that it accepts the
-    Cholesky factor of the correlation matrix to compute probabilities.
+    Reference: Bayesian Data Analysis, 3rd ed., Gelman et al., p. 576
     """
-
-    support = constraints.lower_cholesky
+    support = constraints.positive_definite
 
     def log_prob(self, X):
         if any(s != self.n for s in X.shape[-2:]):
-            raise ValueError("Cholesky factor is not of size n={}".format(self.n.item()))
-        if not _is_valid_correlation_matrix_cholesky_factor(X):
-            raise ValueError("Input is not a Cholesky factor of a valid correlation matrix")
-        log_diag_sum = torch.diagonal(X, dim1=-2, dim2=-1).log().sum(-1)
-        return self.C + (self.eta - 1) * 2 * log_diag_sum
+            raise ValueError("Correlation matrix is not of size n={}".format(self.n))
+        if not _is_valid_correlation_matrix(X):
+            raise ValueError("Input is not a valid correlation matrix")
+        X_cholesky = psd_safe_cholesky(X, upper=False)
+        return super().log_prob(X_cholesky)
+
+    def sample(self, sample_shape=torch.Size()):
+        R = super().sample(sample_shape=sample_shape)
+        return R.matmul(R.transpose(-1, -2))
 
 
 class LKJCovariancePrior(LKJPrior):
@@ -107,15 +83,15 @@ def __init__(self, n, eta, sd_prior, validate_args=False):
             raise ValueError("sd_prior must be an instance of Prior")
         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])}:
+        # bug-fix event shapes if necessary
+        if sd_prior.event_shape not in {torch.Size([1]), torch.Size([n]), torch.Size([])}:
+            if sd_prior.event_shape == torch.Size([]):
+                sd_prior._event_shape = torch.Size([1])
             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=False
-        )
         self.correlation_prior = correlation_prior
         self.sd_prior = sd_prior
 
@@ -130,6 +106,14 @@ def log_prob(self, X):
         log_prob_sd = self.sd_prior.log_prob(marginal_sd)
         return log_prob_corr + log_prob_sd
 
+    def sample(self, sample_shape=torch.Size()):
+        base_correlation = self.correlation_prior.sample(sample_shape)
+        marginal_sds = self.sd_prior.rsample(sample_shape)
+        # expand sds to have the same shape as the base correlation matrix
+        marginal_sds = marginal_sds.repeat(*[1] * len(sample_shape), self.correlation_prior.n)
+        marginal_sds = torch.diag_embed(marginal_sds)
+        return marginal_sds.matmul(base_correlation).matmul(marginal_sds)
+
 
 def _batch_form_diag(tsr):
     """Form diagonal matrices in batch mode."""


Test patch:
diff --git a/test/priors/test_lkj_prior.py b/test/priors/test_lkj_prior.py
--- a/test/priors/test_lkj_prior.py
+++ b/test/priors/test_lkj_prior.py
@@ -4,8 +4,10 @@
 from math import exp
 
 import torch
+from torch.distributions import LKJCholesky
 
 from gpytorch.priors import LKJCholeskyFactorPrior, LKJCovariancePrior, LKJPrior, SmoothedBoxPrior
+from gpytorch.priors.lkj_prior import _is_valid_correlation_matrix, _is_valid_correlation_matrix_cholesky_factor
 from gpytorch.test.utils import approx_equal, least_used_cuda_device
 
 
@@ -14,7 +16,6 @@ def test_lkj_prior_to_gpu(self):
         if torch.cuda.is_available():
             prior = LKJPrior(2, 1.0).cuda()
             self.assertEqual(prior.eta.device.type, "cuda")
-            self.assertEqual(prior.C.device.type, "cuda")
 
     def test_lkj_prior_validate_args(self):
         LKJPrior(2, 1.0, validate_args=True)
@@ -26,18 +27,17 @@ def test_lkj_prior_validate_args(self):
     def test_lkj_prior_log_prob(self, cuda=False):
         device = torch.device("cuda") if cuda else torch.device("cpu")
         prior = LKJPrior(2, torch.tensor(0.5, device=device))
+        dist = LKJCholesky(2, torch.tensor(0.5, device=device))
 
         S = torch.eye(2, device=device)
-        self.assertAlmostEqual(prior.log_prob(S).item(), -1.86942, places=4)
+        S_chol = torch.linalg.cholesky(S)
+        self.assertAlmostEqual(prior.log_prob(S), dist.log_prob(S_chol), 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([-1.86942, -1.72558], device=S.device)))
+        S_chol = torch.linalg.cholesky(S)
+        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
         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 = LKJPrior(2, torch.tensor(1.0, device=device))
-        self.assertTrue(torch.all(prior.log_prob(S) == prior.C))
-
     def test_lkj_prior_log_prob_cuda(self):
         if torch.cuda.is_available():
             with least_used_cuda_device():
@@ -46,11 +46,14 @@ def test_lkj_prior_log_prob_cuda(self):
     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))
+        dist = LKJCholesky(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_chol = torch.linalg.cholesky(S)
+        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
         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)))
+        S_chol = torch.linalg.cholesky(S)
+        self.assertTrue(approx_equal(prior.log_prob(S), dist.log_prob(S_chol)))
         with self.assertRaises(ValueError):
             prior.log_prob(torch.eye(3, device=device))
 
@@ -59,6 +62,18 @@ def test_lkj_prior_batch_log_prob_cuda(self):
             with least_used_cuda_device():
                 self.test_lkj_prior_batch_log_prob(cuda=True)
 
+    def test_lkj_prior_sample(self, seed=0):
+        torch.random.manual_seed(seed)
+
+        prior = LKJPrior(n=5, eta=0.5)
+        random_samples = prior.sample(torch.Size((8,)))
+        self.assertTrue(_is_valid_correlation_matrix(random_samples))
+
+        max_non_symm = (random_samples - random_samples.transpose(-1, -2)).abs().max()
+        self.assertLess(max_non_symm, 1e-4)
+
+        self.assertEqual(random_samples.shape, torch.Size((8, 5, 5)))
+
 
 class TestLKJCholeskyFactorPrior(unittest.TestCase):
     def test_lkj_cholesky_factor_prior_to_gpu(self):
@@ -77,18 +92,13 @@ def test_lkj_cholesky_factor_prior_validate_args(self):
     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))
+        dist = LKJCholesky(2, torch.tensor(0.5, device=device))
         S = torch.eye(2, device=device)
         S_chol = torch.linalg.cholesky(S)
-        self.assertAlmostEqual(prior.log_prob(S_chol).item(), -1.86942, places=4)
+        self.assertAlmostEqual(prior.log_prob(S_chol), dist.log_prob(S_chol), places=4)
         S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S_chol.device)])
         S_chol = torch.stack([torch.linalg.cholesky(Si) 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))
+        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
 
     def test_lkj_cholesky_factor_prior_log_prob_cuda(self):
         if torch.cuda.is_available():
@@ -98,21 +108,26 @@ def test_lkj_cholesky_factor_prior_log_prob_cuda(self):
     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))
+        dist = LKJCholesky(2, torch.tensor([0.5, 1.5], device=device))
 
         S = torch.eye(2, device=device)
         S_chol = torch.linalg.cholesky(S)
-        self.assertTrue(approx_equal(prior.log_prob(S_chol), torch.tensor([-1.86942, -0.483129], device=S_chol.device)))
+        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
         S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
         S_chol = torch.stack([torch.linalg.cholesky(Si) 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))
+        self.assertTrue(approx_equal(prior.log_prob(S_chol), dist.log_prob(S_chol)))
 
     def test_lkj_cholesky_factor_prior_batch_log_prob_cuda(self):
         if torch.cuda.is_available():
             with least_used_cuda_device():
                 self.test_lkj_cholesky_factor_prior_batch_log_prob(cuda=True)
 
+    def test_lkj_prior_sample(self):
+        prior = LKJCholeskyFactorPrior(2, 0.5)
+        random_samples = prior.sample(torch.Size((6,)))
+        self.assertTrue(_is_valid_correlation_matrix_cholesky_factor(random_samples))
+        self.assertEqual(random_samples.shape, torch.Size((6, 2, 2)))
+
 
 class TestLKJCovariancePrior(unittest.TestCase):
     def test_lkj_covariance_prior_to_gpu(self):
@@ -138,17 +153,15 @@ def test_lkj_covariance_prior_log_prob(self, cuda=False):
             sd_prior = sd_prior.cuda()
         prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior)
         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))
+        corr_dist = LKJCholesky(2, torch.tensor(0.5, device=device))
+        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diag()).sum()
+        self.assertAlmostEqual(prior.log_prob(S), dist_log_prob, places=4)
+
+        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
+        S_chol = torch.linalg.cholesky(S)
+        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
+        self.assertTrue(approx_equal(prior.log_prob(S), dist_log_prob))
 
     def test_lkj_covariance_prior_log_prob_cuda(self):
         if torch.cuda.is_available():
@@ -161,18 +174,16 @@ def test_lkj_covariance_prior_log_prob_hetsd(self, cuda=False):
         b = torch.tensor([exp(1), exp(2)], device=device)
         sd_prior = SmoothedBoxPrior(a, b)
         prior = LKJCovariancePrior(2, torch.tensor(0.5, device=device), sd_prior)
+        corr_dist = LKJCholesky(2, torch.tensor(0.5, device=device))
+
         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))
+        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diag()).sum()
+        self.assertAlmostEqual(prior.log_prob(S), dist_log_prob, places=4)
 
-        # 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))
+        S = torch.stack([S, torch.tensor([[1.0, 0.5], [0.5, 1]], device=S.device)])
+        S_chol = torch.linalg.cholesky(S)
+        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
+        self.assertTrue(approx_equal(prior.log_prob(S), dist_log_prob))
 
     def test_lkj_covariance_prior_log_prob_hetsd_cuda(self):
         if torch.cuda.is_available():
@@ -184,19 +195,34 @@ def test_lkj_covariance_prior_batch_log_prob(self, cuda=False):
         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)
+        corr_dist = LKJCholesky(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([-3.59981, -2.21351], device=S.device)))
+        dist_log_prob = corr_dist.log_prob(S) + sd_prior.log_prob(S.diag())
+        self.assertLessEqual((prior.log_prob(S) - dist_log_prob).abs().sum(), 1e-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, -2.35735], device=S.device)))
-        with self.assertRaises(ValueError):
-            prior.log_prob(torch.eye(3, device=device))
+        S_chol = torch.linalg.cholesky(S)
+        dist_log_prob = corr_dist.log_prob(S_chol) + sd_prior.log_prob(torch.diagonal(S, dim1=-2, dim2=-1))
+        self.assertLessEqual((prior.log_prob(S) - dist_log_prob).abs().sum(), 1e-4)
 
     def test_lkj_covariance_prior_batch_log_prob_cuda(self):
         if torch.cuda.is_available():
             with least_used_cuda_device():
                 self.test_lkj_covariance_prior_batch_log_prob(cuda=True)
 
+    def test_lkj_prior_sample(self):
+        prior = LKJCovariancePrior(2, 0.5, sd_prior=SmoothedBoxPrior(exp(-1), exp(1)))
+        random_samples = prior.sample(torch.Size((6,)))
+        # need to check that these are positive semi-sefinite
+        min_eval = torch.linalg.eigh(random_samples)[0].min()
+        self.assertTrue(min_eval >= 0)
+        # and that they are symmetric
+        max_non_symm = (random_samples - random_samples.transpose(-1, -2)).abs().max()
+        self.assertLess(max_non_symm, 1e-4)
+
+        self.assertEqual(random_samples.shape, torch.Size((6, 2, 2)))
+
 
 if __name__ == "__main__":
     unittest.main()