979b8c9efa551e8c948a4aca145367a2d87ac8d6,test/distributions/test_multitask_multivariate_normal.py,TestMultiTaskMultivariateNormal,test_multitask_multivariate_normal,#TestMultiTaskMultivariateNormal#Any#,39

Before Change

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

After Change

    def test_multitask_multivariate_normal(self, cuda=False):
        device = torch.device("cuda") if cuda else torch.device("cpu")
        for dtype in (torch.float, torch.double):
            mean = torch.tensor([[0, 1], [2, 3]], dtype=dtype, device=device)
            variance = torch.tensor([[1, 2], [3, 4]], dtype=dtype, device=device)

            // interleaved
            covmat = variance.view(-1).diag()
            mtmvn = MultitaskMultivariateNormal(mean=mean, covariance_matrix=covmat)
            self.assertTrue(torch.equal(mtmvn.mean, mean))
            self.assertTrue(torch.allclose(mtmvn.variance, variance))
            self.assertTrue(torch.allclose(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, dtype=dtype)).item(), -7.3064, places=4
            )
            logprob = mtmvn.log_prob(torch.zeros(3, 2, 2, device=device, dtype=dtype))
            logprob_expected = -7.3064 * torch.ones(3, device=device, dtype=dtype)
            self.assertTrue(torch.allclose(logprob, logprob_expected))
            conf_lower, conf_upper = mtmvn.confidence_region()
            self.assertTrue(torch.allclose(conf_lower, mtmvn.mean - 2 * mtmvn.stddev))
            self.assertTrue(torch.allclose(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]))

            // non-interleaved
            covmat = variance.transpose(-1, -2).reshape(-1).diag()
            mtmvn = MultitaskMultivariateNormal(mean=mean, covariance_matrix=covmat, interleaved=False)
            self.assertTrue(torch.equal(mtmvn.mean, mean))
            self.assertTrue(torch.allclose(mtmvn.variance, variance))
            self.assertTrue(torch.allclose(mtmvn.scale_tril, covmat.sqrt()))
            self.assertTrue(mtmvn.event_shape == torch.Size([2, 2]))

    def test_multitask_multivariate_normal_cuda(self):
        if torch.cuda.is_available():
            with least_used_cuda_device():
                self.test_multitask_multivariate_normal(cuda=True)

Instances