def test_zip_0_gate(rate):
# if gate is 0 ZIP is Poisson
zip_ = ZeroInflatedPoisson(torch.zeros(1), torch.tensor(rate))
pois = Poisson(torch.tensor(rate))
s = pois.sample((20, ))
zip_prob = zip_.log_prob(s)
pois_prob = pois.log_prob(s)
assert_tensors_equal(zip_prob, pois_prob)
def test_zip_0_gate(rate):
# if gate is 0 ZIP is Poisson
zip1 = ZeroInflatedPoisson(torch.tensor(rate), gate=torch.zeros(1))
zip2 = ZeroInflatedPoisson(torch.tensor(rate),
gate_logits=torch.tensor(-99.9))
pois = Poisson(torch.tensor(rate))
s = pois.sample((20, ))
zip1_prob = zip1.log_prob(s)
zip2_prob = zip2.log_prob(s)
pois_prob = pois.log_prob(s)
assert_close(zip1_prob, pois_prob)
assert_close(zip2_prob, pois_prob)
