//  [python]
def test_multivariate_normal_shape(self):
        mean = torch.randn(5, 3, requires_grad=True)
        mean_no_batch = torch.randn(3, requires_grad=True)
        mean_multi_batch = torch.randn(6, 5, 3, requires_grad=True)

        # construct PSD covariance
        tmp = torch.randn(3, 10)
        cov = (torch.matmul(tmp, tmp.t()) / tmp.size(-1)).requires_grad_()
        prec = cov.inverse().requires_grad_()
        scale_tril = torch.linalg.cholesky(cov, upper=False).requires_grad_()

        # construct batch of PSD covariances
        tmp = torch.randn(6, 5, 3, 10)
        cov_batched = (tmp.unsqueeze(-2) * tmp.unsqueeze(-3)).mean(-1).requires_grad_()
        prec_batched = cov_batched.inverse()
        scale_tril_batched = torch.linalg.cholesky(cov_batched, upper=False)

        # ensure that sample, batch, event shapes all handled correctly
        self.assertEqual(MultivariateNormal(mean, cov).sample().size(), (5, 3))
        self.assertEqual(MultivariateNormal(mean_no_batch, cov).sample().size(), (3,))
        self.assertEqual(MultivariateNormal(mean_multi_batch, cov).sample().size(), (6, 5, 3))
        self.assertEqual(MultivariateNormal(mean, cov).sample((2,)).size(), (2, 5, 3))
        self.assertEqual(MultivariateNormal(mean_no_batch, cov).sample((2,)).size(), (2, 3))
        self.assertEqual(MultivariateNormal(mean_multi_batch, cov).sample((2,)).size(), (2, 6, 5, 3))
        self.assertEqual(MultivariateNormal(mean, cov).sample((2, 7)).size(), (2, 7, 5, 3))
        self.assertEqual(MultivariateNormal(mean_no_batch, cov).sample((2, 7)).size(), (2, 7, 3))
        self.assertEqual(MultivariateNormal(mean_multi_batch, cov).sample((2, 7)).size(), (2, 7, 6, 5, 3))
        self.assertEqual(MultivariateNormal(mean, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))
        self.assertEqual(MultivariateNormal(mean_no_batch, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))
        self.assertEqual(MultivariateNormal(mean_multi_batch, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))
        self.assertEqual(MultivariateNormal(mean, precision_matrix=prec).sample((2, 7)).size(), (2, 7, 5, 3))
        self.assertEqual(MultivariateNormal(mean, precision_matrix=prec_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))
        self.assertEqual(MultivariateNormal(mean, scale_tril=scale_tril).sample((2, 7)).size(), (2, 7, 5, 3))
        self.assertEqual(MultivariateNormal(mean, scale_tril=scale_tril_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3))

        # check gradients
        # We write a custom gradcheck function to maintain the symmetry
        # of the perturbed covariances and their inverses (precision)
        def multivariate_normal_log_prob_gradcheck(mean, covariance=None, precision=None, scale_tril=None):
            mvn_samples = MultivariateNormal(mean, covariance, precision, scale_tril).sample().requires_grad_()

            def gradcheck_func(samples, mu, sigma, prec, scale_tril):
                if sigma is not None:
                    sigma = 0.5 * (sigma + sigma.transpose(-1, -2))  # Ensure symmetry of covariance
                if prec is not None:
                    prec = 0.5 * (prec + prec.transpose(-1, -2))  # Ensure symmetry of precision
                return MultivariateNormal(mu, sigma, prec, scale_tril).log_prob(samples)
            gradcheck(gradcheck_func, (mvn_samples, mean, covariance, precision, scale_tril), raise_exception=True)

        multivariate_normal_log_prob_gradcheck(mean, cov)
        multivariate_normal_log_prob_gradcheck(mean_multi_batch, cov)
        multivariate_normal_log_prob_gradcheck(mean_multi_batch, cov_batched)
        multivariate_normal_log_prob_gradcheck(mean, None, prec)
        multivariate_normal_log_prob_gradcheck(mean_no_batch, None, prec_batched)
        multivariate_normal_log_prob_gradcheck(mean, None, None, scale_tril)
        multivariate_normal_log_prob_gradcheck(mean_no_batch, None, None, scale_tril_batched)