jax

6156decd - [indexing] support pallas-like dynamic slicing in jax.numpy indexing