Repository: jax-ml/jax
Issue category: Performance Issue

Issue:
`jax.random.beta` 3 orders of magnitude slower from 0.4.36 on GPU
### Description

My code runs substantially slower from one month ago, and I figued out a key bottleneck: sampling from beta distribution has gotten around 1000 times slower on GPU.

On Colab, I run the following code on different versions of jax
```
@jax.jit
def sample_beta(rng_key):
  return jax.random.beta(key=rng_key, a=1, b=1, shape=(1000, 2))

seed = jrand.PRNGKey(1)
sample_beta(seed)

%timeit sample_beta(seed)
```

* Time take on version 0.4.35: **0.784ms**
* Time take on version 0.4.36: **351ms**
* Time take on version 0.4.37: **354ms**


![image](https://github.com/user-attachments/assets/e7dd4600-192c-4308-b4d4-52eb7ca762a3)

![image](https://github.com/user-attachments/assets/0e8f6952-50b2-4ba5-8fda-6b3533a94b86)

![image](https://github.com/user-attachments/assets/ce088de7-89e4-466a-b331-3a79cdf77ae2)



### System info (python version, jaxlib version, accelerator, etc.)

jax:    0.4.36
jaxlib: 0.4.36
numpy:  1.26.4
python: 3.10.12 (main, Nov  6 2024, 20:22:13) [GCC 11.4.0]
device info: Tesla T4-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='d36852658d94', release='6.1.85+', version='#1 SMP PREEMPT_DYNAMIC Thu Jun 27 21:05:47 UTC 2024', machine='x86_64')


$ nvidia-smi
Fri Dec 13 13:13:50 2024       
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                       Off | 00000000:00:04.0 Off |                    0 |
| N/A   72C    P0              31W /  70W |    109MiB / 15360MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
+---------------------------------------------------------------------------------------+


Hints:
I can reproduce this, but I'm not totally sure where this would be coming from. Perhaps @jakevdp or @froystig could take a look re: recent changes to PRNGs?
My bisection points to https://github.com/jax-ml/jax/pull/24593

Patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
index 13c4ab4dbce4..12aa5b93efbf 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -291,15 +291,18 @@ def split(key: ArrayLike, num: int | tuple[int, ...] = 2) -> Array:
   return _return_prng_keys(wrapped, _split(typed_key, num))
 
 
-def _key_impl(keys: Array) -> str | PRNGSpec:
+def _key_impl(keys: Array) -> PRNGImpl:
   assert jnp.issubdtype(keys.dtype, dtypes.prng_key)
   keys_dtype = typing.cast(prng.KeyTy, keys.dtype)
-  impl = keys_dtype._impl
+  return keys_dtype._impl
+
+def _key_spec(keys: Array) -> str | PRNGSpec:
+  impl = _key_impl(keys)
   return impl.name if impl.name in prng.prngs else PRNGSpec(impl)
 
 def key_impl(keys: ArrayLike) -> str | PRNGSpec:
   typed_keys, _ = _check_prng_key("key_impl", keys, allow_batched=True)
-  return _key_impl(typed_keys)
+  return _key_spec(typed_keys)
 
 
 def _key_data(keys: Array) -> Array:


Test patch: