Repository: google/jax

Issue:
Efficiently shardable PRNG algorithm for GPU
Please:

- [x] Check for duplicate requests.
- [x] Describe your goal, and if possible provide a code snippet with a motivating example.

https://jax.readthedocs.io/en/latest/jax.random.html#advanced-rng-configuration

In here, the efficiently shardable through jit is only for TPUs. Could you make it for GPU?


Discussion:
Thanks for raising this!

Actually, those docs do _not_ mean to say that the efficiently shardable options, namely `rbg` and `unsafe_rbg`, only work on TPU; they work on CPU and GPU also!

Instead, those docs mean to say that:
* `rbg` and `unsafe_rbg` are the absolute fastest way to generate random values on TPUs, compared to other options on TPUs, and
* the values produced by `rbg` and `unsafe_rbg` on TPU are different than what they are on CPU, even given the same key value, and those both are also different from the values that are produced on GPU. (Perhaps the remarkable thing is that the default RNG, `Threefry`, produces the same raw random bits given the same key on CPU, GPU, and TPU, though numerical floating point differences could lead to downstream differences in values coming out of float-producing random samplers.)

There's another efficiently shardable option which I just noticed isn't listed there, but it's described [here](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#generating-random-numbers). Basically, if you set the option `jax_threefry_partitionable` then the default Threefry PRNG becomes efficiently partitionable. It basically checks all the boxes except "Fastest on TPU", since nothing is as fast as `rbg` and `unsafe_rbg` on TPU.

Does that make sense?

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -106,17 +106,18 @@
 .. table::
    :widths: auto
 
-   =================================   ========  ===  ==========  =======  ==============
-   Property                            Threefry  rbg  unsafe_rbg  rbg (*)  unsafe_rbg (*)
-   =================================   ========  ===  ==========  =======  ==============
-   Fastest on TPU                                ✅   ✅          ✅       ✅
-   efficiently shardable (w/ pjit)                                ✅       ✅
-   identical across shardings           ✅       ✅   ✅
-   identical across CPU/GPU/TPU         ✅
-   identical across JAX/XLA versions    ✅
-   =================================   ========  ===  ==========  =======  ==============
-
-(*): with XLA_FLAGS=xla_tpu_spmd_rng_bit_generator_unsafe=1 set
+   =================================   ========  =========  ===  ==========  =====  ============
+   Property                            Threefry  Threefry*  rbg  unsafe_rbg  rbg**  unsafe_rbg**
+   =================================   ========  =========  ===  ==========  =====  ============
+   Fastest on TPU                                           ✅   ✅          ✅     ✅
+   efficiently shardable (w/ pjit)                ✅                         ✅     ✅
+   identical across shardings           ✅        ✅        ✅   ✅
+   identical across CPU/GPU/TPU         ✅        ✅
+   identical across JAX/XLA versions    ✅        ✅
+   =================================   ========  =========  ===  ==========  =====  ============
+
+(*): with jax_threefry_partitionable=1 set
+(**): with XLA_FLAGS=--xla_tpu_spmd_rng_bit_generator_unsafe=1 set
 
 The difference between "rbg" and "unsafe_rbg" is that while "rbg" uses a less
 robust/studied hash function for random value generation (but not for
@@ -124,6 +125,9 @@
 robust hash functions for `jax.random.split` and `jax.random.fold_in`. Therefore
 less safe in the sense that the quality of random streams it generates from
 different keys is less well understood.
+
+For more about jax_threefry_partitionable, see
+https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html#generating-random-numbers
 """
 
 from jax._src.prng import PRNGKeyArray as _PRNGKeyArray


Test patch: