[JAX] Update users of `jax.ops.index...` functions, which are deprecated.

[copybara-service[bot]]

[JAX] Update users of jax.ops.index... functions, which are deprecated.

• replace uses of jax.ops.index[...] with jax.numpy.index_exp[...], which is a standard NumPy function that does the same thing.
• remove some redundant uses of jax.ops.index[...], where the expression is passed directly to an indexed accessor function like .at[...].
• update some remaining users of jax.ops.index_update(x, jax.ops.index[idx], y) to use the x.at[idx].set(y) APIs.