Repository: google/jax

Issue:
Can jax.random.PRNGKey work with jax_transfer_guard?
It seems that PRNGKey isn't compatible with jax_transfer_guard? I'm getting a `Disallowed host-to-device transfer` error regardless of whether I'm passing in a np or jnp seed.

- Is there a way to use PRNGKey together with jax_transfer_guard? Otherwise, would this be feasible to add?
- As a workaround, can I just replace PRNGKey by creating an array of shape (2,) that contains two uniform random int32 values myself on the outside, or are there more restrictions on what PRNGKey returns?

Example 1:

```python3
import jax
import jax.numpy as jnp

jax.config.update('jax_transfer_guard', 'disallow')

seed = jax.device_put(42)
key = jax.random.PRNGKey(seed)
```

<details>

```
Traceback (most recent call last):
  File "/Users/danijar/temp/example.py", line 7, in <module>
    key = jax.random.PRNGKey(seed)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/random.py", line 128, in PRNGKey
    key = prng.seed_with_impl(impl, seed)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/prng.py", line 262, in seed_with_impl
    return random_seed(seed, impl=impl)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/prng.py", line 555, in random_seed
    return random_seed_p.bind(seeds_arr, impl=impl)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 325, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 328, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 686, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/prng.py", line 567, in random_seed_impl
    base_arr = random_seed_impl_base(seeds, impl=impl)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/prng.py", line 572, in random_seed_impl_base
    return seed(seeds)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/prng.py", line 807, in threefry_seed
    lax.shift_right_logical(seed, lax_internal._const(seed, 32)))
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/lax/lax.py", line 500, in shift_right_logical
    return shift_right_logical_p.bind(x, y)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 325, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 328, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 686, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 113, in apply_primitive
    return compiled_fun(*args)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 198, in <lambda>
    return lambda *args, **kw: compiled(*args, **kw)[0]
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 825, in _execute_compiled
    in_flat = flatten(device_put(x, device) for i, x in enumerate(args)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/util.py", line 113, in concatenate
    return list(it.chain.from_iterable(xs))
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 825, in <genexpr>
    in_flat = flatten(device_put(x, device) for i, x in enumerate(args)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 1115, in device_put
    return device_put_handlers[type(x)](x, device)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 1126, in _device_put_array
    return (backend.buffer_from_pyval(x, device),)
jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: Disallowed host-to-device transfer: type=<class 'numpy.ndarray'>, shape=(), dtype=int32, dst_device=TFRT_CPU_0
```

</details>

Example 2:

```python3
import jax
import jax.numpy as jnp

jax.config.update('jax_transfer_guard', 'disallow')

seed = 42
key = jax.random.PRNGKey(seed)
```

<details>

```
Traceback (most recent call last):
  File "/Users/danijar/temp/example.py", line 7, in <module>
    key = jax.random.PRNGKey(seed)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/random.py", line 128, in PRNGKey
    key = prng.seed_with_impl(impl, seed)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/prng.py", line 262, in seed_with_impl
    return random_seed(seed, impl=impl)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/prng.py", line 552, in random_seed
    seeds_arr = jnp.asarray(np.int64(seeds))
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 1921, in asarray
    return array(a, dtype=dtype, copy=False, order=order)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 1902, in array
    out = lax_internal._convert_element_type(out, dtype, weak_type=weak_type)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/lax/lax.py", line 579, in _convert_element_type
    return convert_element_type_p.bind(operand, new_dtype=new_dtype,
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 325, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 328, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/core.py", line 686, in process_primitive
    return primitive.impl(*tracers, **params)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 113, in apply_primitive
    return compiled_fun(*args)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 198, in <lambda>
    return lambda *args, **kw: compiled(*args, **kw)[0]
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 881, in _execute_trivial
    return [_copy_device_array_to_device(x, device) if device_array.type_is_device_array(x)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 882, in <listcomp>
    else h(None, *device_put(x, device)) for h, x in zip(handlers, outs)]
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 1115, in device_put
    return device_put_handlers[type(x)](x, device)
  File "/Users/danijar/homebrew/lib/python3.9/site-packages/jax/_src/dispatch.py", line 1126, in _device_put_array
    return (backend.buffer_from_pyval(x, device),)
jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: Disallowed host-to-device transfer: type=<class 'numpy.ndarray'>, shape=(), dtype=int32, dst_device=TFRT_CPU_0
```

</details>


Discussion:

Resolving patch:
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -875,6 +875,10 @@ def threefry_seed(seed: typing.Array) -> typing.Array:
     bit-casting to a pair of uint32 values (or from a 32-bit seed by
     first padding out with zeros).
   """
+  return _threefry_seed(seed)
+
+@partial(jit, inline=True)
+def _threefry_seed(seed: typing.Array) -> typing.Array:
   if seed.shape:
     raise TypeError(f"PRNG key seed must be a scalar; got {seed!r}.")
   if not np.issubdtype(seed.dtype, np.integer):


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -527,6 +527,13 @@ def test_isinstance(self):
     self.assertIsInstance(key, jax.Array)
 
 
+class ThreefryPrngTest(jtu.JaxTestCase):
+  def test_seed_no_implicit_transfers(self):
+    # See https://github.com/google/jax/issues/15613
+    with jax.transfer_guard('disallow'):
+      random.threefry2x32_key(jax.device_put(42))  # doesn't crash
+
+
 class LaxRandomTest(jtu.JaxTestCase):
 
   def _CheckCollisions(self, samples, nbits):