Repository: google/jax
Issue: jax.random.poisson does not broadcast to shape of rate param

jax.random.poisson does not broadcast to the shape of the rate parameter. This is a problem for me--I need to JIT compile a function which samples from a Poisson distribution--but because JAX doesn't broadcast the shape of the rate parameter, I end up passing in the shape of the rate parameter to Poisson, and this gives a  `TracerArrayConversionError`.

By contrast, np.random.poisson broadcasts to the shape of the rate parameter.


```python
import jax.numpy as jnp
key = jnp.random.PRNGKey(0)
rate = jnp.ones((5,5,5))
jax.random.poisson(key, lam=rate)
```

Output:
```
~/.conda/envs/jax_env/lib/python3.9/site-packages/jax/_src/random.py in poisson(key, lam, shape, dtype)
   1242   shape = core.canonicalize_shape(shape)
   1243   if np.shape(lam) != shape:
-> 1244     lam = jnp.broadcast_to(lam, shape)
   1245   lam = lax.convert_element_type(lam, np.float32)
   1246   return _poisson(key, lam, shape, dtype)

~/.conda/envs/jax_env/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py in broadcast_to(arr, shape)
   1830     shape_tail = shape[nlead:]
   1831     compatible = _all(core.symbolic_equal_one_of_dim(arr_d, [1, shape_d])
-> 1832                       for arr_d, shape_d in safe_zip(arr_shape, shape_tail))
   1833     if nlead < 0 or not compatible:
   1834       msg = "Incompatible shapes for broadcasting: {} and requested shape {}"

~/.conda/envs/jax_env/lib/python3.9/site-packages/jax/_src/util.py in safe_zip(*args)
     31   n = len(args[0])
     32   for arg in args[1:]:
---> 33     assert len(arg) == n, 'length mismatch: {}'.format(list(map(len, args)))
     34   return list(zip(*args))
     35 

AssertionError: length mismatch: [3, 0]
```

tagging @shoyer since I think he wrote the Poisson sampling code, hope that's not too obnoxious :) 
Link to original PR which added Poisson sampler is here: https://github.com/google/jax/pull/2805

Changed files:
jax/_src/random.py, tests/random_test.py

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 30671c803..5ccb4251b 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1109,20 +1109,21 @@ def _poisson(key, lam, shape, dtype):
 
 def poisson(key: KeyArray,
             lam: RealArray,
-            shape: Sequence[int] = (),
+            shape: Optional[Sequence[int]] = None,
             dtype: DTypeLikeInt = dtypes.int_) -> jnp.ndarray:
   """Sample Poisson random values with given shape and integer dtype.
 
   Args:
     key: a PRNG key used as the random key.
-    lam: rate parameter (mean of the distribution), must be >= 0.
+    lam: rate parameter (mean of the distribution), must be >= 0. Must be broadcast-compatible with ``shape``
     shape: optional, a tuple of nonnegative integers representing the result
-      shape. Default ().
+      shape. Default (None) produces a result shape equal to ``lam.shape``.
     dtype: optional, a integer dtype for the returned values (default int64 if
       jax_enable_x64 is true, otherwise int32).
 
   Returns:
-    A random array with the specified shape and dtype.
+    A random array with the specified dtype and with shape given by ``shape`` if
+    ``shape is not None, or else by ``lam.shape``.
   """
   key, _ = _check_prng_key(key)
   if key.impl is not prng.threefry_prng_impl:
@@ -1130,9 +1131,11 @@ def poisson(key: KeyArray,
         '`poisson` is only implemented for the threefry2x32 RNG, '
         f'not {key.impl}')
   dtype = dtypes.canonicalize_dtype(dtype)
-  shape = core.canonicalize_shape(shape)
-  if np.shape(lam) != shape:
-    lam = jnp.broadcast_to(lam, shape)
+  if shape is not None:
+    shape = core.canonicalize_shape(shape)
+  else:
+    shape = np.shape(lam)
+  lam = jnp.broadcast_to(lam, shape)
   lam = lax.convert_element_type(lam, np.float32)
   return _poisson(key, lam, shape, dtype)
 
diff --git a/tests/random_test.py b/tests/random_test.py
index e5a9898f7..e98dabba1 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -730,6 +730,12 @@ class LaxRandomTest(jtu.JaxTestCase):
     self._CheckChiSquared(samples[:10000], scipy.stats.poisson(2.0).pmf)
     self._CheckChiSquared(samples[10000:], scipy.stats.poisson(20.0).pmf)
 
+  def testPoissonWithoutShape(self):
+    key = self.seed_prng(1)
+    lam = 2 * jnp.ones(10000)
+    samples = random.poisson(key, lam)
+    self._CheckChiSquared(samples, scipy.stats.poisson(2.0).pmf)
+
   def testPoissonShape(self):
     key = self.seed_prng(0)
     x = random.poisson(key, np.array([2.0, 20.0]), shape=(3, 2))