How Jax use LAX-backend implementation of functions

I need to compute the kron procuts of two arrays and I want to test if doing it using Jax is faster than doing it using Numpy.

Now, in numpy my code there is res = numpy.kron(x1,x2), in Jax there is jax.numpy.kron(x1,x2) but how can I use it properly?
My doubs are:

• is it sufficient to replace numpy with jax.numpy as follows: res = jax.numpy.kron(x1,x2)?

• should I first sent x1 and x2 to the device using x1_dev = jax.device_put(x1) and after that run res = jax.numpy.kron(x1_dev,x2_dev)?

• should I add jax.block_until_ready() to the jax.numpy.kron() call?

This is covered in JAX's FAQ under Benchmarking JAX Code.

In particular, if you're interested in the speed of JAX vs NumPy I would read the second section of this, Is JAX Faster Than Numpy? which gives a broad overview of when you should expect JAX to be faster or slower than equivalent NumPy code.

As for benchmarking kron: following the advice there, I would benchmark them like this (using IPython's %timeit for convenience). I ran the following on a Colab T4 GPU runtime:

import numpy as np
x1 = np.random.rand(1000)
x2 = np.random.rand(1000)
%timeit np.kron(x1, x2)

6.52 ms ± 580 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

import jax.numpy as jnp
x1_jax = jnp.array(x1)
x2_jax = jnp.array(x2)
%timeit jnp.kron(x1_jax, x2_jax).block_until_ready()

1.39 ms ± 2.28 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

If you want to see the effect of just-in-time compilation with jax.jit, you can do something like this:

import jax
jit_kron = jax.jit(jnp.kron)
_ = jit_kron(x1_jax, x2_jax)  # trigger compilation before timing
%timeit jit_kron(x1_jax, x2_jax).block_until_ready()

116 µs ± 33.9 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)