<p dir="auto">I have encountered some performance issues with JAX version 0.2.26 and 0.2.27. This has surfaced while running the TestU01 small-crush RNG testsuite against the <code class="notranslate">jax.random</code> RNG as well as my own implementation in <a href="https://github.com/DPBayes/jax-chacha-prng">https://github.com/DPBayes/jax-chacha-prng</a> which run significantly slower than in previous versions. Running JAX version 0.2.26 with JAXlib version 0.1.74 took me 30:47.50 minutes (close to that with JAX version 0.2.27). With JAX version 0.2.25, same JAXlib version, took only 0:38.59 minutes.</p>
<p dir="auto">I'm not sure what causes this and it's hard to track down. It could of course be an issue in my wrapping code around the Python rngs as I haven't yet found a simpler working example that demonstrates this behaviour. I'm trying to narrow it down further, but wanted to ask whether anyone else has encountered something similar or if there were any changes in JAX version 0.2.26 that I'm not aware of and that could explain this.</p>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox"> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:<br>
As said above, I have not boiled this down to a simple complete example.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp
f = lambda x: jnp.add(x, 1)
x = jax.pmap(f)(jnp.array([1]))  # x is a SDA

# jnp.add will use the JAX C++ jit path which has SDA handling falling back to numpy-like handling
x = jnp.add(x, 1)"><pre class="notranslate"><code class="notranslate">import jax
import jax.numpy as jnp
f = lambda x: jnp.add(x, 1)
x = jax.pmap(f)(jnp.array([1]))  # x is a SDA

# jnp.add will use the JAX C++ jit path which has SDA handling falling back to numpy-like handling
x = jnp.add(x, 1)
</code></pre></div>
<p dir="auto">Handling: <a href="https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/python/py_values.cc#L220-L223">https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/python/py_values.cc#L220-L223</a></p>