Issue 1:
This led to a significant performance problem in the threefry2x32 implementation when it was changed to use a dynamic_slice as part of its implementation and when batched as part of the gamma implementation. If we cannot prove that i is equal across batch elements, x[i] turns into a relatively expensive gather of scalars, but if we know x and i are unbatched, it remains a dynamic slice.

Issue 2:
The PR #1175 greatly improves performance of functions involve many random.split ops. Follow up the discussion there about a regression of gamma sampler, I make a repro script here

from jax import jit, lax, random, vmap
import jax.numpy as np
from jax.config import config; config.update("jax_platform_name", "cpu")

def gamma_one(key, alpha):
key, subkey = random.split(key)
boost = np.where(alpha >= 1.0, 1.0, random.uniform(subkey, ()) ** (1.0 / alpha))
alpha = np.where(alpha >= 1.0, alpha, alpha + 1.0)

d = alpha - 1.0 / 3.0
c = 1.0 / np.sqrt(9.0 * d)

def _cond_fn(kXVU):
_, X, V, U = kXVU
return (U >= 1.0 - 0.0331 * X * X) & (np.log(U) >= 0.5 * X + d * (1.0 - V + np.log(V)))

def _body_fn(kXVU):
def _next_kxv(kxv):
k = kxv[0]
k, subkey = random.split(k)
x = random.normal(subkey, ())
v = 1.0 + c * x
return k, x, v

key = kXVU[0]
key, x_key, U_key = random.split(key, 3)
_, x, v = lax.while_loop(lambda kxv: kxv[2] <= 0.0, _next_kxv, (x_key, 0.0, -1.0))
X = x * x
V = v * v * v
U = random.uniform(U_key, ())
return key, X, V, U

_, _, V, _ = lax.while_loop(_cond_fn, _body_fn, (key, 1.0, 1.0, 2.0))
z = d * V * boost
return np.where(z == 0, np.finfo(z.dtype).tiny, z)

def gamma_trick_one(key, alpha):
boost = np.where(alpha >= 1.0, 1.0, random.uniform(key, ()) ** (1.0 / alpha))
key, = random.split(key, 1)
alpha = np.where(alpha >= 1.0, alpha, alpha + 1.0)

d = alpha - 1.0 / 3.0
c = 1.0 / np.sqrt(9.0 * d)

def _cond_fn(kXVU):
_, X, V, U = kXVU
return (U >= 1.0 - 0.0331 * X * X) & (np.log(U) >= 0.5 * X + d * (1.0 - V + np.log(V)))

def _body_fn(kXVU):
def _next_kxv(kxv):
k = kxv[0]
x = random.normal(k, ())
k, = random.split(k, 1)
v = 1.0 + c * x
return k, x, v

key = kXVU[0]
key, x, v = lax.while_loop(lambda kxv: kxv[2] <= 0.0, _next_kxv, (key, 0.0, -1.0))
X = x * x
V = v * v * v
U = random.uniform(key, ())
key, = random.split(key, 1)
return key, X, V, U

_, _, V, _ = lax.while_loop(_cond_fn, _body_fn, (key, 1.0, 1.0, 2.0))
z = d * V * boost
return np.where(z == 0, np.finfo(z.dtype).tiny, z)

gamma_trick = jit(vmap(gamma_trick_one))
gamma = jit(vmap(gamma_one))
rng = random.PRNGKey(0)
alpha = random.uniform(rng, (1000,)) * 2
rngs = random.split(rng, 1000)

then

%time x1 = gamma(rngs, alpha).copy()
%time x2 = gamma_trick(rngs, alpha).copy()
%time x3 = random.gamma(rng, alpha).copy()
assert (np.abs(x1 - x3) > 1e-5).sum() == 0

returns (with JAX 0.1.41)

CPU times: user 2.06 s, sys: 28 ms, total: 2.09 s
Wall time: 2.07 s
CPU times: user 1.41 s, sys: 7.77 ms, total: 1.42 s
Wall time: 1.41 s
CPU times: user 6.01 s, sys: 35.8 ms, total: 6.04 s
Wall time: 6.03 s

To my surprise, time to compile gamma and random.gamma are so different, while the algorithm is the same (modulo precision issues due to using np. vs lax. - I don't understand why precision is 1e-5 but I guess that it is due to how constant values are handled in two implementations). Maybe we should use gamma implementation instead of the current implementation in random.gamma?

In current master branch (after PR #1175), the compiling time of them are

CPU times: user 4.43 s, sys: 11.1 ms, total: 4.45 s
Wall time: 4.42 s
CPU times: user 4.17 s, sys: 31.9 ms, total: 4.21 s
Wall time: 4.2 s
CPU times: user 4.79 s, sys: 16.2 ms, total: 4.8 s
Wall time: 4.78 s

which improves the performance of random.gamma but makes regression for gamma and gamma_trick.

@mattjj I hope the above repro code illustrates some interesting performance issue for JAX/XLA team. :)