jnp.argsort much slower than the numpy version

Here’s a comparison of the JAX and numpy versions of argsort on a CPU:

import numpy as np
import jax.numpy as jnp
from jax import config, random
config.update('jax_platform_name', 'cpu')

key = random.PRNGKey(42)
key, subkey = random.split(key)

x_jnp = random.uniform(subkey, (100, 10000))
x_np = np.array(x_jnp)

%%timeit
np.argsort(x_np, axis=0)

%%timeit
jnp.argsort(x_jnp, axis=0).block_until_ready()

In this case jnp.argsort is ~5X slower than than np.argsort. I’m seeing >20x difference with more realistic arrays. Why is there such a large difference in performance between the two implementations?

[jakevdp]

You might find this FAQ helpful: FAQ: Is JAX Faster Than NumPy?.

[hawkinsp]

I also note that the slowness is specific to floating-point values. Sorting int32 values is significantly faster. The only difference between the two as far as I can tell is the comparison function.