Jax indexing code runs extremely slowly?

[arjun-mani]

This code takes more than 15 seconds to run, while running it in standard numpy is instantaneous. I've looked around but can't find any explanation why. Intuitively, it feels like doing a simple indexing operation should not take nearly this long. Would appreciate any insight here, thank you so much.

import jax
import jax.numpy as jnp
from jax import jit

@jit
def foo(a):
	for i in range(128):
		for j in range(128):
			s = a[0]
	return 0

 
a = jnp.zeros(128)
foo(a)

(128 is arbitrary, just happens to be relevant to my use-case).

[wrzadkow]

Code compiled using jax.jit generally takes a lot of time when run for the first time, because the just-in-time compilation is performed then. All subsequent runs should be much faster.

I tried your code and replaced your last line with

t = time.time()
foo(a)
print(f'Including just-in-time compilation {time.time()-t}')
t = time.time()
foo(a)
print(f'Run using compiled code {time.time()-t}')

and got on my setup (time is measured in seconds)

Including just-in-time compilation 35.65647530555725
Run using compiled code 0.0005266666412353516

[machineko]

I meet the same problem at first with my code, so I've check common gotcha's especially here in-place-updates

You can see that the first call will probably create new array every time you assign value to array element and will copy it (it was case with my code :P)

Also, first JIT call can be a lot longer cause of unrolling of the loops inside for_loops

You can easily fix long compilation time running it on smaller array at first and using lax loops and some time just using numpy for looping and variable assignment is just a lot faster. (In my case almost all my code based on creating new arrays and padding just use pure numpy and then convert it to jax arrrays it's a lot faster when you got a lot of dynamic shaped arrays)

@jit
def foo(a):
	for i in range(128):
		for j in range(128):
			s = a[0]  # here you probably create new array every time assigment happend without jitted function in pure python it'll be a 'pointer'
	return 0

[GaoyuanWu]

Hi!

In my case almost all my code based on creating new arrays and padding just use pure numpy and then convert it to jax arrrays it's a lot faster when you got a lot of dynamic shaped arrays

Do you mean something like this:

def f(x):
  arr = np.zeros(128)
  for i in range(128):
    arr[i] = assign_something(i,x) #function with i & x inputs
  jax_arr = jnp.array(arr)
  return jax_arr

This should be a lot faster if x is static.
But what if x is traced? Is there a "fast way"? (Since arr[i] = assign_something(i,x) will give an error if I am not mistaken)
Thanks a lot!

[jakevdp]

There shouldn't be any problem with traced inputs in this code – if you could open a new discussion thread with a more complete example of what you're trying to do and the problem you're facing, that would be a better place to get an answer than a comment on a two-year-old thread. Thanks!

[jakevdp]

The reason this takes a long time to run is because you are timing the tracing and compiling of the function. Once the compilation is finished, subsequent runs will be fast:

import jax
import jax.numpy as jnp
from jax import jit

@jit
def foo(a):
  for i in range(32):
    for j in range(32):
      s = a[0]
  return 0

 
a = jnp.zeros(128)
%time foo(a).block_until_ready()  # trace, compile & run
# CPU times: user 2.38 s, sys: 28.2 ms, total: 2.41 s
# Wall time: 2.4 s
%time foo(a).block_until_ready()  # run the compiled function
# CPU times: user 1.03 ms, sys: 0 ns, total: 1.03 ms
# Wall time: 827 µs

Why is compilation so slow here? Well, jax.jit effectively unrolls all Python loops in the course of tracing the function, so it essentially turns this code snippet into a flat list of 128 * 128=16384 operations in sequence, and then this list of operations is fed to XLA for compilation. In practice, I believe XLA compilation scales as roughly the square of the number of operations, so code like this performs rather poorly.

You can read more about how JIT interacts with Python control flow in https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#python-control-flow-+-JIT