JAX seems slow in basic computation?

[Emerald01]

At Colab, I am doing the following experiment to see the speed of some basic addition computation, the code is very simple

sum and a conditional sum in JAX

N = jax.numpy
L = jax.lax

x = N.array([1,2,3,4,5,6,7,8,9,10])
y = N.array([1,2,3,4,5,6,7,8,9,10])

@jit
def jax_sum(x, y):
  return x + y

@jit
def jax_cond_sum(x,y):  
  z = L.cond(N.all(x>y),
         lambda op: op[0] - op[1],     
         lambda op: op[0] + op[1],
         (x, y))
  return z

same way, in Numpy

x = np.array([1,2,3,4,5,6,7,8,9,10])
y = np.array([1,2,3,4,5,6,7,8,9,10])

def np_sum(x, y):
  return x + y

def np_cond_sum(x,y):
  if np.all(x > y):
    return x+y
  else:
    return x-y

%timeit -n 10 -r 10 jax_sum(x, y)
%timeit -n 10 -r 10 jax_cond_sum(x, y)

For CPU only

10 loops, best of 10: 87 µs per loop
10 loops, best of 10: 54.2 µs per loop

For GPU accelerator

10 loops, best of 10: 49.8 µs per loop
10 loops, best of 10: 67.6 µs per loop

%timeit -n 10 -r 10 np_sum(x, y)
%timeit -n 10 -r 10 np_cond_sum(x, y)

10 loops, best of 10: 692 ns per loop
10 loops, best of 10: 5.03 µs per loop

It seems that a plain numpy has orders of magnitude win, why is that, I already used @jit and the code is straightforward to run, do I make anything wrong? Where can we get the gain using JAX

[jakevdp]

I can't run this, because I don't know what L is.

In general, though, JAX has a fixed-cost overhead for every computation that is on the order tens of milliseconds, while numpy does not. I suspect that if you ran your benchmarks on arrays that are large enough so that this overhead is not important, you'd see a truer comparison of JAX vs numpy runtime scaling.

Also, keep in mind for this sort of micro-benchmark that JAX executes Asynchronously, and so you should use the .block_until_ready() method on any outputs for which you want accurate computation times.