<commit_before>from __future__ import absolute_import, division, print_function, unicode_literals

import unittest
from math import exp

import torch
from gpytorch.priors import LKJCholeskyFactorPrior, LKJCovariancePrior, LKJPrior, SmoothedBoxPrior
from test._utils import approx_equal


class TestLKJPrior(unittest.TestCase):
    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)
        with self.assertRaises(ValueError):
            LKJPrior(1.5, 1.0, validate_args=True)
        with self.assertRaises(ValueError):
            LKJPrior(2, -1.0, validate_args=True)

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

        self.assertFalse(prior.log_transform)
        S = torch.eye(2, device=device)
        self.assertAlmostEqual(prior.log_prob(S).item(), -1.86942, 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)))
        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():
            return self.test_lkj_prior_log_prob(cuda=True)

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

        self.assertFalse(prior.log_transform)
        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_lkj_prior_batch_log_prob_cuda(self):
        if torch.cuda.is_available():
            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.cholesky(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.cholesky(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.cholesky(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.cholesky(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):<commit_msg>Remove needless upper=False for cholesky<commit_after>from __future__ import absolute_import, division, print_function, unicode_literals

import unittest
from math import exp

import torch
from gpytorch.priors import LKJCholeskyFactorPrior, LKJCovariancePrior, LKJPrior, SmoothedBoxPrior
from test._utils import approx_equal


class TestLKJPrior(unittest.TestCase):
    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)
        with self.assertRaises(ValueError):
            LKJPrior(1.5, 1.0, validate_args=True)
        with self.assertRaises(ValueError):
            LKJPrior(2, -1.0, validate_args=True)

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

        self.assertFalse(prior.log_transform)
        S = torch.eye(2, device=device)
        self.assertAlmostEqual(prior.log_prob(S).item(), -1.86942, 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)))
        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():
            return self.test_lkj_prior_log_prob(cuda=True)

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

        self.assertFalse(prior.log_transform)
        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_lkj_prior_batch_log_prob_cuda(self):
        if torch.cuda.is_available():
            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.cholesky(S)
        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.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))

    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.cholesky(S)
        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.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))

    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):
{'fraction_non_alphanumeric': 0.13051048674317373, 'fraction_numerical': 0.029442026117926394, 'mean_word_length': 3.4983980064079745, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2736611', 'n_tokens_mistral': 4814, 'n_tokens_neox': 4445, 'n_words': 663}