Lower more index_update cases to slice

[jakevdp]

Before:

In [1]: import jax.numpy as jnp
In [2]: import jax
In [3]: x = jnp.arange(24).reshape(2, 3, 4)
In [4]: jax.make_jaxpr(lambda x: x[1,:,:2])(x)
Out[4]: 
{ lambda ; a:i32[2,3,4]. let
    b:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 1
    c:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
    d:i32[2] = concatenate[dimension=0] b c
    e:i32[3,2] = gather[
      dimension_numbers=GatherDimensionNumbers(offset_dims=(0, 1), collapsed_slice_dims=(0,), start_index_map=(0, 2))
      fill_value=None
      indices_are_sorted=True
      mode=GatherScatterMode.PROMISE_IN_BOUNDS
      slice_sizes=(1, 3, 2)
      unique_indices=True
    ] a d
  in (e,) }

After:

In [4]: jax.make_jaxpr(lambda x: x[1,:,:2])(x)
Out[4]: 
{ lambda ; a:i32[2,3,4]. let
    b:i32[1,3,2] = slice[
      limit_indices=(2, 3, 2)
      start_indices=(1, 0, 0)
      strides=(1, 1, 1)
    ] a
    c:i32[3,2] = squeeze[dimensions=(0,)] b
  in (c,) }

[yashk2810]

Instead of just checking the jaxpr, please check the output too

See this function: https://github.com/google/jax/blob/main/tests/array_test.py#L481-L484

[yashk2810]

Also please test the slicing with Arrays sharded over multiple devices

You can add more cases here if you want: https://github.com/google/jax/blob/main/tests/array_test.py#L480

[jakevdp]

From offline discussion: it looks like lax.slice is fundamentally broken for sharded arrays, so @yashk2810 suggested that a way forward here is probably to skip the new code path if the input array is sharded.

[jakevdp]

closing in favor of #15377