//  [python]
def test_wishart_shape(self):
        set_rng_seed(0)  # see Note [Randomized statistical tests]
        ndim = 3

        df = torch.rand(5, requires_grad=True) + ndim
        df_no_batch = torch.rand([], requires_grad=True) + ndim
        df_multi_batch = torch.rand(6, 5, requires_grad=True) + ndim

        # construct PSD covariance
        tmp = torch.randn(ndim, 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, ndim, 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, ndim, ndim))
        self.assertEqual(Wishart(df_no_batch, cov).sample().size(), (ndim, ndim))
        self.assertEqual(
            Wishart(df_multi_batch, cov).sample().size(), (6, 5, ndim, ndim)
        )
        self.assertEqual(Wishart(df, cov).sample((2,)).size(), (2, 5, ndim, ndim))
        self.assertEqual(Wishart(df_no_batch, cov).sample((2,)).size(), (2, ndim, ndim))
        self.assertEqual(
            Wishart(df_multi_batch, cov).sample((2,)).size(), (2, 6, 5, ndim, ndim)
        )
        self.assertEqual(Wishart(df, cov).sample((2, 7)).size(), (2, 7, 5, ndim, ndim))
        self.assertEqual(
            Wishart(df_no_batch, cov).sample((2, 7)).size(), (2, 7, ndim, ndim)
        )
        self.assertEqual(
            Wishart(df_multi_batch, cov).sample((2, 7)).size(), (2, 7, 6, 5, ndim, ndim)
        )
        self.assertEqual(
            Wishart(df, cov_batched).sample((2, 7)).size(), (2, 7, 6, 5, ndim, ndim)
        )
        self.assertEqual(
            Wishart(df_no_batch, cov_batched).sample((2, 7)).size(),
            (2, 7, 6, 5, ndim, ndim),
        )
        self.assertEqual(
            Wishart(df_multi_batch, cov_batched).sample((2, 7)).size(),
            (2, 7, 6, 5, ndim, ndim),
        )
        self.assertEqual(
            Wishart(df, precision_matrix=prec).sample((2, 7)).size(),
            (2, 7, 5, ndim, ndim),
        )
        self.assertEqual(
            Wishart(df, precision_matrix=prec_batched).sample((2, 7)).size(),
            (2, 7, 6, 5, ndim, ndim),
        )
        self.assertEqual(
            Wishart(df, scale_tril=scale_tril).sample((2, 7)).size(),
            (2, 7, 5, ndim, ndim),
        )
        self.assertEqual(
            Wishart(df, scale_tril=scale_tril_batched).sample((2, 7)).size(),
            (2, 7, 6, 5, ndim, ndim),
        )

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