def METHOD_NAME(
    mvn_dist, k, sample_shape, L21, omega1, L11, L22=0.8, L33=0.9, omega2=0.75
):
    if mvn_dist == "OMTMultivariateNormal" and k > 1:
        return
    omega = torch.tensor([omega1, omega2, 0.0])
    loc = torch.zeros(3, requires_grad=True)
    zero_vec = [0.0, 0.0, 0.0]
    off_diag = torch.tensor([zero_vec, [L21, 0.0, 0.0], zero_vec], requires_grad=True)
    L = torch.diag(torch.tensor([L11, L22, L33])) + off_diag
    if mvn_dist == "OMTMultivariateNormal":
        dist = OMTMultivariateNormal(loc, L)
    elif mvn_dist == "AVFMultivariateNormal":
        CV = (1.1 * torch.rand(2, k, 3)).requires_grad_(True)
        dist = AVFMultivariateNormal(loc, L, CV)
    z = dist.rsample(sample_shape)
    torch.cos((omega * z).sum(-1)).mean().backward()
    computed_grad = off_diag.grad.cpu().data.numpy()[1, 0]
    analytic = analytic_grad(L11=L11, L22=L22, L21=L21, omega1=omega1, omega2=omega2)
    assert off_diag.grad.size() == off_diag.size()
    assert loc.grad.size() == loc.size()
    assert torch.triu(off_diag.grad, 1).sum() == 0.0
    assert_equal(
        analytic,
        computed_grad,
        prec=0.005,
        msg="bad cholesky grad for %s (expected %.5f, got %.5f)"
        % (mvn_dist, analytic, computed_grad),
    )