Repository: google/jax

Issue:
compiling time grows non-linear w.r.t. the number of random.split ops
A repro code
```
from jax import jit, random
from jax.config import config; config.update("jax_platform_name", "cpu")
import time

def f(key, i):
    for _ in range(i):
        _, key = random.split(key)
    return key

for i in range(17):
    t = time.time()
    key = jit(f, static_argnums=(1,))(random.PRNGKey(0), i).copy()
    print("split {} times takes {}".format(i, time.time() - t))
```
which returns
```
split 0 times takes 0.008060216903686523
split 1 times takes 0.05638694763183594
split 2 times takes 0.13771629333496094
split 3 times takes 0.284271240234375
split 4 times takes 0.4720790386199951
split 5 times takes 0.8609561920166016
split 6 times takes 1.5761287212371826
split 7 times takes 2.615913152694702
split 8 times takes 4.20055079460144
split 9 times takes 6.581374168395996
split 10 times takes 9.83530855178833
split 11 times takes 15.45705771446228
split 12 times takes 20.33790898323059
split 13 times takes 27.418986320495605
split 14 times takes 35.04498267173767
split 15 times takes 44.33582162857056
split 16 times takes 53.186384439468384
```

@mattjj I think that this causes the regression of gamma sampler in our last discussion. The above issue is causing a big problem in compiling models with many random latent variables in NumPyro.


Discussion:
Thanks for opening this. We'll look into it.

Can you send me a master list of JAX bugs that are blocking or otherwise inconveniencing your work? I want to prioritize them, but I fear that I've lost track of the most important ones.
This isn't a great workaround, but just to make sure you're aware, if you make just one `split` call things are faster, i.e. `all_keys = random.split(key, 16)`.

I looked at a profile and found that all the time is being spent in thew XLA compiler. Each call to `random.split` is inlining a call to the hash function underlying `random.split`, and that inlining means that we're building large XLA graphs (and XLA compile times can be super-linear in the input size). There may be things we can do on the JAX side, and also I pinged the XLA folks for ideas.
Thank Matt, that's a great idea! We'll think more if we can go with that solution. cc @neerajprad

> Can you send me a master list of JAX bugs that are blocking or otherwise inconveniencing your work?

Currently, beside this issue, we only have some problems with fastmath mode, which makes it impossible (due to `nan` issue) for us to do stochastic/mcmc inference on some models involving e.g. logsumexp. I'll resort some repro code and will make a separate issue for it soon. Other than that, JAX <3 works beautifully for us.

Resolving patch:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -32,7 +32,7 @@
 from . import numpy as np
 from . import tree_util
 from .api import custom_transforms, defjvp, jit, vmap
-from .numpy.lax_numpy import _constant_like, asarray
+from .numpy.lax_numpy import _constant_like, asarray, stack
 from jax.lib import xla_bridge
 from jax import core
 from jax.scipy.special import logit
@@ -123,38 +123,22 @@ def apply_round(v, rot):
   else:
     x = list(np.split(count.ravel(), 2))
 
-  rotations = onp.uint32([13, 15, 26, 6, 17, 29, 16, 24])
-  ks = [key1, key2, key1 ^ key2 ^ onp.uint32(0x1BD11BDA)]
+  rotations = onp.array([[13, 15, 26, 6], [17, 29, 16, 24]], dtype=onp.uint32)
+  ks = stack([key1, key2, key1 ^ key2 ^ onp.uint32(0x1BD11BDA)])
+  idxs = onp.array([[1, 2], [2, 0], [0, 1], [1, 2], [2, 0]])
 
   x[0] = x[0] + ks[0]
   x[1] = x[1] + ks[1]
 
-  for r in rotations[:4]:
-    x = apply_round(x, r)
-  x[0] = x[0] + ks[1]
-  x[1] = x[1] + ks[2] + onp.uint32(1)
+  get = partial(lax.dynamic_index_in_dim, keepdims=False)
 
-  for r in rotations[4:]:
-    x = apply_round(x, r)
-  x[0] = x[0] + ks[2]
-  x[1] = x[1] + ks[0] + onp.uint32(2)
-
-  for r in rotations[:4]:
-    x = apply_round(x, r)
-  x[0] = x[0] + ks[0]
-  x[1] = x[1] + ks[1] + onp.uint32(3)
-
-  for r in rotations[4:]:
-    x = apply_round(x, r)
-  x[0] = x[0] + ks[1]
-  x[1] = x[1] + ks[2] + onp.uint32(4)
-
-  for r in rotations[:4]:
-    x = apply_round(x, r)
-  x[0] = x[0] + ks[2]
-  x[1] = x[1] + ks[0] + onp.uint32(5)
-
-  out = np.concatenate(x)
+  def step(i, x):
+    for r in get(rotations, i % 2):
+      x = apply_round(x, r)
+    i0, i1 = get(idxs, i)
+    return [x[0] + ks[i0],
+            x[1] + ks[i1] + asarray(i + 1, dtype=onp.uint32)]
+  out = np.concatenate(lax.fori_loop(0, 5, step, x))
   assert out.dtype == onp.uint32
   return lax.reshape(out[:-1] if odd_size else out, count.shape)
 


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -398,10 +398,12 @@ def testIssue756(self):
       self.assertEqual(onp.result_type(w), onp.float32)
 
   def testNoOpByOpUnderHash(self):
-    def fail(): assert False
+    def fail(*args, **kwargs): assert False
     apply_primitive, xla.apply_primitive = xla.apply_primitive, fail
-    out = random.threefry_2x32(onp.zeros(2, onp.uint32), onp.arange(10, dtype=onp.uint32))
-    xla.apply_primitive = apply_primitive
+    try:
+      out = random.threefry_2x32(onp.zeros(2, onp.uint32), onp.arange(10, dtype=onp.uint32))
+    finally:
+      xla.apply_primitive = apply_primitive
 
 
 if __name__ == "__main__":