jax

80f18ded - [sharding_in_types] Make slice and ellipsis work with `.at[...].get(out_sharding=P(...))`