Fix spurious KeyReuseError under vmap over a zero-sized key array

jax/experimental/key_reuse/_core.py

@@ -76,10 +76,12 @@ def __init__(self, idx: int, mask: bool | np.bool_ | np.ndarray = True):
     assert isinstance(idx, int)
     if isinstance(mask, np.ndarray):
       assert mask.dtype == np.dtype('bool')
-      if np.all(mask):
-        mask = True
-      elif not np.any(mask):
+      if not np.any(mask):
+        # An empty mask (e.g. from vmap over a zero-sized axis) is a no-op:
+        # np.all() is vacuously True for empty arrays, so this must come first.
         mask = False
+      elif np.all(mask):
+        mask = True
       elif mask.flags.writeable:
           mask = np.array(mask, copy=True)
           mask.flags.writeable = False

tests/key_reuse_test.py

@@ -673,6 +673,16 @@ def f():
     with self.assertRaisesRegex(KeyReuseError, self.traced_bits_msg):
       f()
 
+  def test_vmap_zero_sized_axis(self):
+    # Regression test for https://github.com/jax-ml/jax/issues/37859: vmap over
+    # a zero-sized key axis must not raise a spurious KeyReuseError.
+    def f(key):
+      a, b = jax.random.split(key)
+      return jax.random.bits(a) + jax.random.bits(b)
+    keys = jax.random.split(jax.random.key(0), 0)
+    self.assertEqual(keys.shape, (0,))
+    jax.vmap(f)(keys)  # does not raise
+
 
 class KeyReuseImplementationTest(jtu.JaxTestCase):