<p dir="auto">This led to a significant performance problem in the <code class="notranslate">threefry2x32</code> implementation when it was changed to use a <code class="notranslate">dynamic_slice</code> as part of its implementation and when batched as part of the <code class="notranslate">gamma</code> implementation. If we cannot prove that <code class="notranslate">i</code> is equal across batch elements, <code class="notranslate">x[i]</code> turns into a relatively expensive gather of scalars, but if we know <code class="notranslate">x</code> and <code class="notranslate">i</code> are unbatched, it remains a dynamic slice.</p>
<p dir="auto">The PR <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="480184913" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/1175" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/1175/hovercard" href="https://github.com/google/jax/pull/1175">#1175</a> greatly improves performance of functions involve many <code class="notranslate">random.split</code> ops. Follow up the discussion there about a regression of gamma sampler, I make a repro script here</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="from jax import jit, lax, random, vmap
import jax.numpy as np
from jax.config import config; config.update(&quot;jax_platform_name&quot;, &quot;cpu&quot;)

def gamma_one(key, alpha):
    key, subkey = random.split(key)
    boost = np.where(alpha &gt;= 1.0, 1.0, random.uniform(subkey, ()) ** (1.0 / alpha))
    alpha = np.where(alpha &gt;= 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 &gt;= 1.0 - 0.0331 * X * X) &amp; (np.log(U) &gt;= 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] &lt;= 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 &gt;= 1.0, 1.0, random.uniform(key, ()) ** (1.0 / alpha))
    key, = random.split(key, 1)
    alpha = np.where(alpha &gt;= 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 &gt;= 1.0 - 0.0331 * X * X) &amp; (np.log(U) &gt;= 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] &lt;= 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)"><pre class="notranslate"><code class="notranslate">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 &gt;= 1.0, 1.0, random.uniform(subkey, ()) ** (1.0 / alpha))
    alpha = np.where(alpha &gt;= 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 &gt;= 1.0 - 0.0331 * X * X) &amp; (np.log(U) &gt;= 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] &lt;= 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 &gt;= 1.0, 1.0, random.uniform(key, ()) ** (1.0 / alpha))
    key, = random.split(key, 1)
    alpha = np.where(alpha &gt;= 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 &gt;= 1.0 - 0.0331 * X * X) &amp; (np.log(U) &gt;= 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] &lt;= 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)
</code></pre></div>
<p dir="auto">then</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="%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) &gt; 1e-5).sum() == 0"><pre class="notranslate"><code class="notranslate">%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) &gt; 1e-5).sum() == 0
</code></pre></div>
<p dir="auto">returns (with JAX 0.1.41)</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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"><pre class="notranslate"><code class="notranslate">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
</code></pre></div>
<p dir="auto">To my surprise, time to compile <code class="notranslate">gamma</code> and <code class="notranslate">random.gamma</code> are so different, while the algorithm is the same (modulo precision issues due to using <code class="notranslate">np.</code> vs <code class="notranslate">lax.</code> - 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 <code class="notranslate">gamma</code> implementation instead of the current implementation in <code class="notranslate">random.gamma</code>?</p>
<p dir="auto">In current master branch (after PR <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="480184913" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/1175" data-hovercard-type="pull_request" data-hovercard-url="/google/jax/pull/1175/hovercard" href="https://github.com/google/jax/pull/1175">#1175</a>), the compiling time of them are</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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"><pre class="notranslate"><code class="notranslate">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
</code></pre></div>
<p dir="auto">which improves the performance of <code class="notranslate">random.gamma</code> but makes regression for <code class="notranslate">gamma</code> and <code class="notranslate">gamma_trick</code>.</p>
<p dir="auto"><a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mattjj/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mattjj">@mattjj</a> I hope the above repro code illustrates some interesting performance issue for JAX/XLA team. :)</p>