Unable to load BART model, AttributeError: module 'jax.ops' has no attribute 'index_update'
## Environment info

- `transformers` version: 4.17.0
- Platform: Linux-5.11.0-1018-gcp-x86_64-with-glibc2.31
- Python version: 3.10.4
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.4.1 (tpu)
- Jax version: 0.3.4
- JaxLib version: 0.3.2
- Using GPU in script?: no
- Using distributed or parallel set-up in script?: no

### Who can help

@patil-suraj @patrickvonplaten

## Information

Model I am using: BART

The problem arises when using:
* [x] the official example scripts

## To reproduce

Steps to reproduce the behavior:

```python
>>> from transformers import FlaxBartForSequenceClassification
>>> model = FlaxBartForSequenceClassification.from_pretrained('facebook/bart-base')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/ayaka/venv310/lib/python3.10/site-packages/transformers/modeling_flax_utils.py", line 550, in from_pretrained
    model = cls(config, *model_args, **model_kwargs)
  File "/home/ayaka/venv310/lib/python3.10/site-packages/transformers/models/bart/modeling_flax_bart.py", line 933, in __init__
    super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
  File "/home/ayaka/venv310/lib/python3.10/site-packages/transformers/modeling_flax_utils.py", line 116, in __init__
    random_params = self.init_weights(self.key, input_shape)
  File "/home/ayaka/venv310/lib/python3.10/site-packages/transformers/models/bart/modeling_flax_bart.py", line 939, in init_weights
    input_ids = jax.ops.index_update(input_ids, (..., -1), self.config.eos_token_id)
AttributeError: module 'jax.ops' has no attribute 'index_update'
```

https://jax.readthedocs.io/en/latest/jax.ops.html#module-jax.ops

> The functions jax.ops.index_update, jax.ops.index_add, etc., which were deprecated in JAX 0.2.22, have been removed. Please use the [jax.numpy.ndarray.at](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at) property on JAX arrays instead.
<|||||>Fixed in #16078