Why lax operations are slower than jnp operations?

[riven314]

From my understanding, the operations from jax.numpy will internally call their associated primitive operations from jax.lax module, so their speed should be comparable.

But I found operations from jax.lax are generally slower than that from jax.numpy (as attached below). Why there is such a discrepancy in speed?

FYI: I am using JAX version 0.2.25 in Colab (CPU mode)

[jakevdp]

I think the reason is that most jnp operations are jit-compiled by default, whereas lax operations are not.

You can see this with:

with jax.disable_jit():
  %timeit lax.sin(1.).block_until_ready()
  %timeit lax.cos(1.).block_until_ready()
  %timeit jnp.sin(1.).block_until_ready()
  %timeit jnp.cos(1.).block_until_ready()

(Note block_until_ready() to account for Asynchronous dispatch)