Speedup T5 Flax training by using Numpy instead of JAX for batch shuffling
# What does this PR do?

The example Flax T5 MLM script was modified to not use JAX for batch shuffling.

With this change, training time per step decreased from 0.246 s to 0.185 s, a speedup of 1.3 (t5-base model trained on TPU v3-8 VM, batch size 64, Adafactor optimizer).

Shuffling on CPU also prevents using memory on the accelerator devices for the batch index array, which is significant for the large datasets that are typically used with pre-training. For instance, the batch index array size for a 30GB dataset with sequence length 512 is 500MB.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR. @patrickvonplaten 
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_15963). All of your documentation changes will be reflected on that endpoint.<|||||>Think I'm in favor of this. However could we also set a numpy random seed in the script to make the runs reproducible? 

cc @patil-suraj - what do you think here? <|||||>> However could we also set a numpy random seed in the script to make the runs reproducible?

The script calls trainer_utils.py's set_seed() which amongst others sets numpy's seed.