def METHOD_NAME(dim):
    dist = LKJCholesky(dim, torch.tensor([1.0]))
    a_sample = dist.sample(torch.Size([100]))
    lp = dist.log_prob(a_sample)
    if dim == 2:
        assert_equal(lp, lp.new_full(lp.size(), -math.log(2)))
    else:
        ladj = (
            a_sample.diagonal(dim1=-2, dim2=-1)
            .log()
            .mul(
                torch.linspace(
                    start=dim - 1,
                    end=0,
                    steps=dim,
                    device=a_sample.device,
                    dtype=a_sample.dtype,
                )
            )
            .sum(-1)
        )
        lps_less_ladj = lp - ladj
        assert (lps_less_ladj - lps_less_ladj.min()).abs().sum() < 1e-4