def METHOD_NAME(total_count, batch_shape, is_sparse):
    event_shape = (3,)
    concentration = torch.rand(batch_shape + event_shape).exp()
    # test on one-hots
    value = total_count * torch.eye(3).reshape(
        event_shape + (1,) * len(batch_shape) + event_shape
    )
    num_samples = 100000
    probs = dist.Dirichlet(concentration).sample((num_samples, 1))
    log_probs = dist.Multinomial(total_count, probs).log_prob(value)
    assert log_probs.shape == (num_samples,) + event_shape + batch_shape
    expected = log_probs.logsumexp(0) - math.log(num_samples)
    actual = DirichletMultinomial(concentration, total_count, is_sparse).log_prob(value)
    assert_close(actual, expected, atol=0.05)