def METHOD_NAME(total_count, probs):
    # if gate is 0 ZINB is NegativeBinomial
    zinb1 = ZeroInflatedNegativeBinomial(
        total_count=torch.tensor(total_count),
        gate=torch.zeros(1),
        probs=torch.tensor(probs),
    )
    zinb2 = ZeroInflatedNegativeBinomial(
        total_count=torch.tensor(total_count),
        gate_logits=torch.tensor(-99.9),
        probs=torch.tensor(probs),
    )
    neg_bin = NegativeBinomial(torch.tensor(total_count), probs=torch.tensor(probs))
    s = neg_bin.sample((20,))
    zinb1_prob = zinb1.log_prob(s)
    zinb2_prob = zinb2.log_prob(s)
    neg_bin_prob = neg_bin.log_prob(s)
    assert_close(zinb1_prob, neg_bin_prob)
    assert_close(zinb2_prob, neg_bin_prob)