//  [python]
def test_wishart_shape(self):
        df = (torch.rand(5, requires_grad=True) + 1) * 10
        df_no_batch = (torch.rand([], requires_grad=True) + 1) * 10
        df_multi_batch = (torch.rand(6, 5, requires_grad=True) + 1) * 10

        # 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.cholesky(cov).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.cholesky(cov_batched)

        # ensure that sample, batch, event shapes all handled correctly
        self.assertEqual(Wishart(df, cov).sample().size(), (5, 3, 3))
        self.assertEqual(Wishart(df_no_batch, cov).sample().size(), (3, 3))
        self.assertEqual(Wishart(df_multi_batch, cov).sample().size(), (6, 5, 3, 3))
        self.assertEqual(Wishart(df, cov).sample((2,)).size(), (2, 5, 3, 3))
        self.assertEqual(Wishart(df_no_batch, cov).sample((2,)).size(), (2, 3, 3))
        self.assertEqual(Wishart(df_multi_batch, cov).sample((2,)).size(), (2, 6, 5, 3, 3))
        self.assertEqual(Wishart(df, cov).sample((2, 7)).size(), (2, 7, 5, 3, 3))
        self.assertEqual(Wishart(df_no_batch, cov).sample((2, 7)).size(), (2, 7, 3, 3))
        self.assertEqual(Wishart(df_multi_batch, cov).sample((2, 7)).size(), (2, 7, 6, 5, 3, 3))
        self.assertEqual(Wishart(df, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3, 3))
        self.assertEqual(Wishart(df_no_batch, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3, 3))
        self.assertEqual(Wishart(df_multi_batch, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3, 3))
        self.assertEqual(Wishart(df, precision_matrix=prec).sample((2, 7)).size(), (2, 7, 5, 3, 3))
        self.assertEqual(Wishart(df, precision_matrix=prec_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3, 3))
        self.assertEqual(Wishart(df, scale_tril=scale_tril).sample((2, 7)).size(), (2, 7, 5, 3, 3))
        self.assertEqual(Wishart(df, scale_tril=scale_tril_batched).sample((2, 7)).size(), (2, 7, 6, 5, 3, 3))

        # check gradients
        # Modified and applied the same tests for multivariate_normal
        def wishart_log_prob_gradcheck(df=None, covariance=None, precision=None, scale_tril=None):
            wishart_samples = Wishart(df, covariance, precision, scale_tril).sample().requires_grad_()

            def gradcheck_func(samples, nu, sigma, prec, scale_tril):
                if sigma is not None:
                    sigma = 0.5 * (sigma + sigma.mT)  # Ensure symmetry of covariance
                if prec is not None:
                    prec = 0.5 * (prec + prec.mT)  # Ensure symmetry of precision
                if scale_tril is not None:
                    scale_tril = scale_tril.tril()
                return Wishart(nu, sigma, prec, scale_tril).log_prob(samples)
            gradcheck(gradcheck_func, (wishart_samples, df, covariance, precision, scale_tril), raise_exception=True)

        wishart_log_prob_gradcheck(df, cov)
        wishart_log_prob_gradcheck(df_multi_batch, cov)
        wishart_log_prob_gradcheck(df_multi_batch, cov_batched)
        wishart_log_prob_gradcheck(df, None, prec)
        wishart_log_prob_gradcheck(df_no_batch, None, prec_batched)
        wishart_log_prob_gradcheck(df, None, None, scale_tril)
        wishart_log_prob_gradcheck(df_no_batch, None, None, scale_tril_batched)