Conversion from a JAX array to a Numpy array is slower after calling a JAX function

[jiayingqi]

Dear JAX community, I found that the conversion from a JAX array to a Numpy array is slower after calling a JAX function. This slowing down persists even if that JAX function has nothing to do with the array conversion. The slowing down scales with the computational complexity of the JAX function.

Results:
After calling the function: the first run is normal, but the remaining ones are slower.

0.00045911199413239956
0.0025374609977006912
0.0026233160169795156
0.0026847750414162874
0.002685735991690308

Without calling the JAX function: everything is normal.

0.0004427499952726066
0.00032009504502639174
0.00032241898588836193
0.00031042599584907293
0.0002709300024434924

This issue appears on both CPU and GPU, and here is an example for reproducing the issue. I would like to ask if there are any fixes.

import time
import jax
import numpy as np

W = jax.random.uniform(jax.random.key(0), (200000, 7))
def compute(x):  # Some JAX function
    return [email protected]

in_data = jax.random.uniform(jax.random.key(0), (1000, 7))
in_data = np.array(in_data)

def test():
    compute(in_data)   # This line slows down the remaining conversion task
    other_data = jax.random.uniform(jax.random.key(0), 30000)
    start = time.perf_counter()
    other_data = np.array(other_data)
    print(time.perf_counter() - start)

for _ in range(5):
    test()

[jakevdp]

I'm not able to reproduce this with JAX v0.4.33 on either a Colab CPU or GPU runtime. That said, I have a guess as to why this may be happening for you: I suspect your timings are being fooled by JAX's Asyncronous dispatch: when you call a JAX operation, the Python function returns before the result is actually computed. So your first call to compute essentially just queues up the computations, which begin running in the background. By the time you get to your second call, the queue is full, and so the Python function must wait for the previous iterations to finish before it can enqueue its computations.

If I'm right, then adding this to your first line should make all runs take the same amount of time:

jax.block_until_ready(compute(in_data))

(This is assuming that compute returns all results that it computes: the block_until_ready will only apply to the values passed to it)

More broadly, if you're interested in microbenchmarks of JAX code, I'd suggest starting here to make sure you're measuring what you think you're measuring: FAQ – benchmarking JAX code.