Repository: jax-ml/jax (python)

Issue:
jax_debug_key_reuse incorrectly raises error
### Description

The following code raises error `KeyReuseError: Previously-consumed key passed to random_split at index 0` despite that we haven't used key2 yet.

```python
import jax; jax.config.update('jax_debug_key_reuse', True)
key = jax.random.key(0)
key1, key2 = jax.random.split(key)
_, key2 = jax.device_put((key1, key2))
jax.random.split(key2)
```

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.36
jaxlib: 0.4.36
numpy:  1.26.4
python: 3.12.7 (main, Oct  1 2024, 02:05:46) [Clang 16.0.0 (clang-1600.0.26.3)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='phandu-mac.roam.internal', release='24.2.0', version='Darwin Kernel Version 24.2.0: Fri Dec  6 19:02:41 PST 2024; root:xnu-11215.61.5~2/RELEASE_ARM64_T6030', machine='arm64')
```

Pull request description:
[key reuse] fix signature for device_put
Fixes #25953

Patch:
diff --git a/jax/experimental/key_reuse/_core.py b/jax/experimental/key_reuse/_core.py
index ef19f94c4..87c41e13c 100644
--- a/jax/experimental/key_reuse/_core.py
+++ b/jax/experimental/key_reuse/_core.py
@@ -292,7 +292,6 @@ key_reuse_signatures[random.random_gamma_p] = KeyReuseSignature(Sink(0))
 key_reuse_signatures[lax.broadcast_in_dim_p] = KeyReuseSignature(Forward(0, 0))
 key_reuse_signatures[lax.copy_p] = KeyReuseSignature(Forward(0, 0))
 key_reuse_signatures[lax.convert_element_type_p] = KeyReuseSignature(Forward(0, 0))
-key_reuse_signatures[lax.device_put_p] = KeyReuseSignature(Forward(0, 0))
 key_reuse_signatures[lax.reshape_p] = KeyReuseSignature(Forward(0, 0))
 key_reuse_signatures[lax.squeeze_p] = KeyReuseSignature(Forward(0, 0))
 key_reuse_signatures[prng.random_wrap_p] = KeyReuseSignature(Source(0))
@@ -561,6 +560,14 @@ def _remat_key_type_signature(eqn):
 key_reuse_signatures[remat_p] = _remat_key_type_signature
 
 
+@dynamic_key_reuse_signature
+def _device_put_signature(eqn):
+  num_vals = len(eqn.invars)
+  return KeyReuseSignature(*(Forward(i, i) for i in range(num_vals)))
+
+key_reuse_signatures[lax.device_put_p] = _device_put_signature
+
+
 def call_impl_with_key_reuse_checks(prim: core.Primitive, raw_impl: Callable[..., Any], *args, **kwargs) -> Any:
   if prim not in key_reuse_signatures:
     # TODO(jakevdp): should we use an unknown signature here?


Test patch:
diff --git a/tests/key_reuse_test.py b/tests/key_reuse_test.py
index 3364c9be9..0d7ac9c18 100644
--- a/tests/key_reuse_test.py
+++ b/tests/key_reuse_test.py
@@ -49,7 +49,6 @@ primitives_with_static_signatures = {
   jax.lax.broadcast_in_dim_p: (lambda key: key[None], key),
   jax.lax.copy_p: (jnp.array, key),
   jax.lax.convert_element_type_p: (lambda key: jnp.array(key, dtype=key.dtype), key),
-  jax.lax.device_put_p: (jax.device_put, key),
   jax.lax.reshape_p: (lambda key: key.reshape((1,)), key),
   jax.lax.squeeze_p: (jnp.squeeze, key1D),
   jax.lax.dynamic_slice_p: (partial(jax.lax.dynamic_slice, slice_sizes=(1,)), key1D, (0,)),
@@ -178,14 +177,27 @@ class KeyReuseUnitTestWithForwarding(jtu.JaxTestCase):
   def test_device_put(self):
     def f(key):
       assert_unconsumed(key)
-      key2 = jax.device_put(key)
-      assert_unconsumed(key)
-      assert_unconsumed(key2)
+      key_d = jax.device_put(key)
+      assert_unconsumed(key_d)
       consume(key)
-      assert_consumed(key)
-      assert_consumed(key2)
+      assert_consumed(key_d)
     self.check_key_reuse(f, jax.random.key(0))
 
+  def test_device_put_multiple(self):
+    def f(key1, key2):
+      assert_unconsumed(key1)
+      assert_unconsumed(key2)
+      key1_d, key2_d = jax.device_put((key1, key2))
+
+      assert_unconsumed(key1_d)
+      consume(key1)
+      assert_consumed(key1_d)
+
+      assert_unconsumed(key2_d)
+      consume(key2)
+      assert_consumed(key2_d)
+    self.check_key_reuse(f, jax.random.key(0), jax.random.key(1))
+
   def test_squeeze(self):
     def f(key):
       assert_unconsumed(key)