Bug in jax.random.poisson for lam=0
Dear people from JAX,

Thank you for writing this software, I love it! I am using JAX for some optics simulations for which I need to use the poissonian random number generator to add noise to my simulated measurements. I noticed that `poisson(key, lam=0)` always returns -1, which makes no sense, and it does not match with the numpy implementation.

In JAX, I would write:

```python
import jax.numpy as jnp
import jax
key = jax.random.PRNGKey(0)
original = jnp.zeros(10,)
x = jax.random.poisson(key,original,(10,))
print(x)
# [-1 -1 -1 -1 -1 -1 -1 -1 -1 -1]
```
Numpy:

```python
import numpy as np
rng = np.random.default_rng()
s = rng.poisson(0, 10)
print(s)
# [0 0 0 0 0 0 0 0 0 0]
```

Looking at the source the problem seems to be in _poisson_knuth. An obvious solution would be to add a statement that just returns zero wherever lam==0,  but there's probably a more elegant solution?
diff --git a/jax/_src/random.py b/jax/_src/random.py
index d6447cd8e..778080067 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1173,11 +1173,12 @@ def _poisson(key, lam, shape, dtype):
   # λ -> ∞, so pick some arbitrary large value.
   lam_rejection = lax.select(use_knuth, lax.full_like(lam, 1e5), lam)
   max_iters = dtype.type(jnp.iinfo(dtype).max)  # insanely conservative
-  return lax.select(
-      use_knuth,
-      _poisson_knuth(key, lam_knuth, shape, dtype, max_iters),
-      _poisson_rejection(key, lam_rejection, shape, dtype, max_iters),
+  result = lax.select(
+    use_knuth,
+    _poisson_knuth(key, lam_knuth, shape, dtype, max_iters),
+    _poisson_rejection(key, lam_rejection, shape, dtype, max_iters),
   )
+  return lax.select(lam == 0, jnp.zeros_like(result), result)
 
 
 def poisson(key, lam, shape=(), dtype=dtypes.int_):
diff --git a/tests/random_test.py b/tests/random_test.py
index d1174c437..3b7dc7e12 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -538,6 +538,12 @@ class LaxRandomTest(jtu.JaxTestCase):
     x = random.poisson(key, np.array([2.0, 20.0]), shape=(3, 2))
     assert x.shape == (3, 2)
 
+  def testPoissonZeros(self):
+    key = random.PRNGKey(0)
+    lam = jnp.concatenate([jnp.zeros(10), 20 * jnp.ones(10)])
+    samples = random.poisson(key, lam, shape=(2, 20))
+    self.assertArraysEqual(samples[:, :10], jnp.zeros_like(samples[:, :10]))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_dtype={}".format(np.dtype(dtype).name), "dtype": dtype}
       for dtype in jtu.dtypes.floating))
['jax/_src/random.py', 'tests/random_test.py']