Use the same device for splitting keys.

PiperOrigin-RevId: 556806725
google-deepmind · Aug 14, 2023 · 3e503c9 · 3e503c9
1 parent 1377a0d
commit 3e503c9
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/haiku/_src/base.py b/haiku/_src/base.py
@@ -1073,7 +1073,8 @@ def reserve(self, num):
       #     key, subkey3, subkey4 = jax.random.split(key, 3)  # reserve(2)
       #
       # Where subkey1->subkey4 are provided to the user in order when requested.
-      new_keys = tuple(jax.random.split(self._key, num + 1))
+      with jax.default_device(self._key.device()):
+        new_keys = tuple(jax.random.split(self._key, num + 1))
       self._key = new_keys[0]
       self._subkeys.extend(new_keys[1:])