Merge pull request #16427 from jakevdp:fix-select-vmap

PiperOrigin-RevId: 540858315
google · Jun 16, 2023 · ba10e94 · ba10e94
2 parents 154c4da + 47ce94f
commit ba10e94
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git 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
@@ -1860,6 +1860,15 @@ def test_select(self):
     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):