Crash when running `examples/flax/question-answering`
Author: DwarKapex
Pull request: no

### System Info

- `transformers` version: 4.36.0.dev0
- Platform: Linux-5.15.0-88-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.19.4
- Safetensors version: 0.4.0
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): not installed (NA)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.8.0 (gpu)
- Jax version: 0.4.21.dev20231121+g2efa5862a
- JaxLib version: 0.4.21.dev20231121
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@sanchit-gandhi @ArthurZucker  @younesbelkada

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. Install the latest version of HF transformer:
```
   $> git clone https://github.com/huggingface/transformers.git /opt/transformers
   $> cd /opt/transformers
   $> pip install -e .
```
2. Navigate to examples:
```
    $> cd /opt/transformers/examples/flax/question-answering
```
3. Install requirements:
```
    $> pip install -r requirements.txt
```
4. Run test:
```
   $> python run_qa.py \
  --model_name_or_path bert-base-uncased \
  --dataset_name squad \
  --do_train   \
  --do_eval   \
  --max_seq_length 384 \
  --doc_stride 128 \
  --learning_rate 3e-5 \
  --num_train_epochs 2 \
  --per_device_train_batch_size 12 \
  --output_dir ./bert-qa-squad \
  --eval_steps 1000
```
5. Crash with the following error:
```
Traceback (most recent call last):
  File "/opt/transformers/examples/flax/question-answering/run_qa.py", line 1095, in <module>
    main()
  File "/opt/transformers/examples/flax/question-answering/run_qa.py", line 900, in main
    model = FlaxAutoModelForQuestionAnswering.from_pretrained(
  File "/opt/transformers/src/transformers/models/auto/auto_factory.py", line 566, in from_pretrained
    return model_class.from_pretrained(
  File "/opt/transformers/src/transformers/modeling_flax_utils.py", line 902, in from_pretrained
    model = cls(config, *model_args, _do_init=_do_init, **model_kwargs)
  File "/opt/transformers/src/transformers/models/bert/modeling_flax_bert.py", line 786, in __init__
    super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)
  File "/opt/transformers/src/transformers/modeling_flax_utils.py", line 219, in __init__
    random_params = self.init_weights(self.key, input_shape)
  File "/opt/transformers/src/transformers/models/bert/modeling_flax_bert.py", line 821, in init_weights
    module_init_outputs = self.module.init(
  File "/opt/transformers/src/transformers/models/bert/modeling_flax_bert.py", line 1572, in __call__
    start_logits, end_logits = logits.split(self.config.num_labels, axis=-1)
AttributeError: 'ArrayImpl' object has no attribute 'split'. Did you mean: '_split'?
```

### Expected behavior

Expect no exception

Comments:
Meet same issue when JAX v0.4.20.