eval_shape appears to hang when tracing a large random sample
Calling

```python
import jax
from jax import random
jax.eval_shape(lambda: random.normal(key=random.PRNGKey(0), shape=[int(1e11)]))
```

returns `ShapeDtypeStruct(shape=(100000000000,), dtype=float32)`. I'd expect that calling

```python
jax.eval_shape(lambda: random.normal(key=random.PRNGKey(0), shape=[int(1e11), int(1e11)]))
```

would similarly return `ShapeDtypeStruct(shape=(100000000000,100000000000), dtype=float32)` (of course, if we *actually* ran the function we'd get an OOM error). 

But instead, this call seems to end up in an ~infinite loop (at least, it hangs for > 2 minutes). Reproduction here:
https://colab.research.google.com/drive/14IZF3WUTaq976jaUC5tehKTPwt3YhNTe?usp=sharing

Is this expected?

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 4fb534ae5..965c4c120 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -304,14 +304,15 @@ def _random_bits(key, bit_width, shape):
   max_count = int(np.ceil(bit_width * size / 32))
 
   nblocks, rem = divmod(max_count, jnp.iinfo(np.uint32).max)
+
   if not nblocks:
     bits = threefry_2x32(key, lax.iota(np.uint32, rem))
   else:
-    *subkeys, last_key = split(key, nblocks + 1)
-    blocks = [threefry_2x32(k, lax.iota(np.uint32, jnp.iinfo(np.uint32).max))
-              for k in subkeys]
+    keys = split(key, nblocks + 1)
+    subkeys, last_key = keys[:-1], keys[-1]
+    blocks = vmap(threefry_2x32, in_axes=(0, None))(subkeys, lax.iota(np.uint32, jnp.iinfo(np.uint32).max))
     last = threefry_2x32(last_key, lax.iota(np.uint32, rem))
-    bits = lax.concatenate(blocks + [last], 0)
+    bits = lax.concatenate([blocks.ravel(), last], 0)
 
   dtype = _UINT_DTYPES[bit_width]
   if bit_width == 64:
['jax/_src/random.py']