Converting batch into jax arrays during training is inefficient
### System Info

The Flax Transformers examples use a custom dataloader which yields batches. This includes converting the batch into jax arrays: 

https://github.com/huggingface/transformers/blob/main/examples/flax/summarization/run_summarization_flax.py#L355

This operation is inefficient, it can take about 700ms to generate a batch during training. On A100 GPUs, this bottleneck leads to lower GPU utilization. 

The solution would be to set ds.set_format("jax") before loading batches for training. 

### Who can help?

@patil-suraj 

### 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

All examples in flax use this operation for batch generation:
https://github.com/huggingface/transformers/blob/main/examples/flax/

### Expected behavior

The fix gives better GPU utilization and faster training (Tested locally on A100 GPUs). 
Hi @isunitha98selvan. In the example linked, we actually force the batch indices to `numpy` arrays:
https://github.com/huggingface/transformers/blob/286a18fa0080dd39bd373008d11d831fbb1a77f1/examples/flax/summarization/run_summarization_flax.py#L345-L346

And then slice the dataset according to these `numpy` arrays.

When using JAX, we want to keep everything that's outside of the jit'd function as a `numpy` array, and everything that's inside as `jax.numpy` arrays. The reason being if everything outside of the jit'd function is a `numpy` array, we can make maximal use of JAX's asynchronous dispatch (https://jax.readthedocs.io/en/latest/async_dispatch.html): the CPU can run ahead and compute any `numpy` arrays ahead of time, whilst the GPUs/TPUs tackle the `jax.numpy` arrays inside the jit'd function.

In the forward pass of the model, any `numpy` arrays are converted to `jax.numpy` arrays, putting the data on the accelerator device.

So whilst it might be slower computing the batches on a per-batch basis like this, when we combine it with a jit'd training function we'll actually train much faster overall 🚀<|||||>Hey @isunitha98selvan, feel free to ask if you have any questions! Otherwise, closing this for now.