Interaction between `vmap` and new PRNG
### Description

Hi, I'm having trouble seeing how the following MWE using [Diffrax](https://github.com/patrick-kidger/diffrax) errors:
```py
import jax
from jax import numpy as jnp
from diffrax import diffeqsolve, ControlTerm, Euler, MultiTerm, ODETerm, SaveAt, VirtualBrownianTree

B = 20
D = 5

tol = 1e-3
dt0 = 0.05
T0 = 0.0
Tf = 1.0

def f(x0, key):
    drift = lambda t, x, args: jnp.zeros((D, ))
    diffusion = lambda t, x, args: jnp.zeros((D, D))
    brownian_motion = VirtualBrownianTree(T0, Tf, tol=tol, shape=x0.shape, key=key)
    terms = MultiTerm(ODETerm(drift), ControlTerm(diffusion, brownian_motion))
    solver = Euler()
    saveat = SaveAt(dense=True)
    sol = diffeqsolve(terms, solver, T0, Tf, dt0=dt0, y0=x0, saveat=saveat)
    return sol.evaluate(Tf)

x0 = jnp.zeros((D, ))
key = jax.random.key(1)
print(f(x0, key))

x0 = jnp.zeros((B, D))
key = jax.random.split(key, B)
print(jax.vmap(f)(x0, key))
```
In particular, the last three lines (with `vmap`) cause an error, while the previous three lines seem to work fine.
I raised [this issue in Diffrax](https://github.com/patrick-kidger/diffrax/issues/267) and was told that it was an issue only with the new PRNG, so the bug is in Jax rather than Diffrax. Hence I'm raising it here.

Here are links for:
- [The stack trace.](https://pastebin.com/46hbkCKb)
- [The output of my `conda list`.](https://pastebin.com/3xjrfYYz)

### What jax/jaxlib version are you using?

0.4.12/0.4.12

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.10.11, WSL 2 -> Ubuntu 22.04.2 LTS

### NVIDIA GPU info

_No response_
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 6151b48af..c397c9fdb 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -3488,7 +3488,6 @@ def _select_hlo_lowering_opaque(ctx, which, *cases):
   avals_in = ctx.avals_in
   aval_out, = ctx.avals_out
   assert all(aval_case == aval_out for aval_case in avals_in[1:])
-  assert avals_in[0].ndim == aval_out.ndim
   select_lower = _select_hlo_lowering
 
   physical_aval_out = core.physical_aval(aval_out)
diff --git a/tests/random_test.py b/tests/random_test.py
index a08ba1b15..c99942519 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1860,6 +1860,15 @@ class KeyArrayTest(jtu.JaxTestCase):
     self.assertIsInstance(ys, random.KeyArray)
     self.assertEqual(ys.shape, (3, 2))
 
+  @skipIf(not config.jax_enable_custom_prng,
+          'requires config.jax_enable_custom_prng')
+  def test_select_scalar_cond(self):
+    # regression test for https://github.com/google/jax/issues/16422
+    ks = self.make_keys(3)
+    ys = lax.select(True, ks, ks)
+    self.assertIsInstance(ys, random.KeyArray)
+    self.assertEqual(ys.shape, (3,))
+
   @skipIf(not config.jax_enable_custom_prng,
           'requires config.jax_enable_custom_prng')
   def test_select2(self):
['jax/_src/lax/lax.py', 'tests/random_test.py']