[Flax] remove jax.ops.index
# What does this PR do?

Remove `jax.ops.index` as this method is removed in latest jax release `v0.3.2`.
_The documentation is not available anymore as the PR was closed or merged._<|||||>There seem to be some failing tests still :-) <|||||>The failing tests are not related to this PR.<|||||>Hi, I am not very into the detail here, but curious about that if this implies from this version, `transformers` will only work with the new version of `jax`? <|||||>> Hi, I am not very into the detail here, but curious about that if this implies from this version, `transformers` will only work with the new version of `jax`?

No, the `.at` syntax is also supported in older jax versions.<|||||>Feel free to merge if you want @patil-suraj 