def test_sampling_shape(lmbda, expected_shape, sample_shape, rng_key):
poisson_samples = Poisson(lmbda=lmbda).sample(rng_key, sample_shape)
assert poisson_samples.shape == expected_shape
def test_logpdf_out_of_support(case):
logprob = Poisson(1.0).logpdf(case["x"])
assert logprob == case["expected"]
def test_sampling_noshape(lmbda, expected_shape, rng_key):
assert Poisson(lmbda=lmbda).sample(rng_key).shape == expected_shape