jax

098e1332 - Error out if the arr is sharded over the axes being rolled over in `jnp.roll`