`replace` argument changes how `jax.random.choice` treats `shape`
Using jax version 0.1.75 on a colab cpu runtime.
```python
import jax

key = jax.random.PRNGKey(0)
jax.random.choice(key, 5, 2, replace=False)
jax.random.choice(key, 5, 2, replace=True)
```
Error:
```
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-19-56696338bc13> in <module>()
      1 key = jax.random.PRNGKey(0)
      2 jax.random.choice(key, 5, 2, replace=False)
----> 3 jax.random.choice(key, 5, 2, replace=True)

2 frames
/usr/local/lib/python3.6/dist-packages/jax/random.py in choice(key, a, shape, replace, p)
    566   if p is None:
    567     if replace:
--> 568       ind = randint(key, shape, 0, n_inputs)
    569       result = ind if a.ndim == 0 else a[ind]
    570     else:

/usr/local/lib/python3.6/dist-packages/jax/random.py in randint(key, shape, minval, maxval, dtype)
    417   """
    418   dtype = dtypes.canonicalize_dtype(dtype)
--> 419   shape = abstract_arrays.canonicalize_shape(shape)
    420   return _randint(key, shape, minval, maxval, dtype)
    421 

/usr/local/lib/python3.6/dist-packages/jax/core.py in canonicalize_shape(shape)
   1078          "got {}.")
   1079   if any(isinstance(x, Tracer) and isinstance(get_aval(x), ShapedArray)
-> 1080          and not isinstance(get_aval(x), ConcreteArray) for x in shape):
   1081     msg += ("\\nIf using `jit`, try using `static_argnums` or applying `jit` to "
   1082             "smaller subfunctions.")

TypeError: 'int' object is not iterable
```
diff --git a/jax/random.py b/jax/random.py
index 985bccae4..2c775d924 100644
--- a/jax/random.py
+++ b/jax/random.py
@@ -556,6 +556,9 @@ def choice(key, a, shape=(), replace=True, p=None):
   Returns:
     An array of shape `shape` containing samples from `a`.
   """
+  if not isinstance(shape, Sequence):
+    raise TypeError("shape argument of jax.random.choice must be a sequence, "
+                    f"got {shape}")
   a = jnp.asarray(a)
   if a.ndim not in [0, 1]:
     raise ValueError("a must be an integer or 1-dimensional")
diff --git a/tests/random_test.py b/tests/random_test.py
index 7151955fe..e665bf0f2 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -848,6 +848,13 @@ class LaxRandomTest(jtu.JaxTestCase):
       self.assertAllClose(
           counts[1]/ num_samples, 0.5, rtol=1e-02, atol=1e-02)
 
+  def testChoiceShapeIsNotSequenceError(self):
+    key = random.PRNGKey(0)
+    with self.assertRaises(TypeError):
+      random.choice(key, 5, 2, replace=False)
+    with self.assertRaises(TypeError):
+      random.choice(key, 5, 2, replace=True)
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
['jax/random.py', 'tests/random_test.py']