jnp.split() has long compile times for large numbers of splits

[jakevdp]

Reported in https://stackoverflow.com/q/74199437

import jax.numpy as jnp
x = jnp.ones(5000)
%time _ = jnp.split(x, jnp.arange(0, 5000, 10))
# CPU times: user 7.29 s, sys: 101 ms, total: 7.39 s
# Wall time: 7.46 s

Subsequent calls are much faster, once the XLA:Slice operations have been cached:

%time _ = jnp.split(x, jnp.arange(0, 5000, 10))
# CPU times: user 61.4 ms, sys: 0 ns, total: 61.4 ms
# Wall time: 70.3 ms

There are some tradeoffs here that make the solution to this non-trivial (see my answer in the above StackOverflow post), but it might be better to use lax.dynamic_slice in place of lax.slice to prevent this kind of catastrophically-bad outcome.

[jakevdp]

Closing as duplicate of #9445