Unexpected slow-down with @jit on simple functions that only use element-wise operations and jnp.roll() on CPUs

[pmocz]

See issue raised with JAX: jax-ml/jax#24373

Simple colab reproducer: https://colab.research.google.com/drive/1pnnC3-XY8-K0MvUv9gH_on2-USG1ZaeY?usp=sharing

Unexpected major slowdown particular to the XLA:CPU compiler is seen when using @jax.jit on a fairly simple function that only use element-wise operations and jnp.roll() . Works well running on TPUs