Make data shuffling in `run_clm_flax.py` respect global seed
Author: bminixhofer
Pull request: yes

# What does this PR do?

Use `jax.random.permutation` instead of `np.random.permutation` in the `data_loader` function of `run_clm_flax.py` to make it use the global seed. Currently batch order would probably vary across runs, regardless of the global seed.

Also changes `np.arange` to `jnp.arange` and `np.array` to `jnp.array` and removes the numpy import, although that would not be strictly necessary.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] 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?

@patil-suraj


Comments:
Thanks a lot for the PR @bminixhofer, good catch!

However, initially, we actually used `jax.random.permutation` and `jax.arange` in the data loader but then switched to numpy as we observed it was causing some issues with JAX's asynchronous dispatch. Since JAX by default puts everything on the device it could cause some issues (especially on TPU) if the dataloader /collator is used with multiple threads to do background fetching. This also leads to major slowdowns. So all flax examples now don't use JAX functions in pre-processing, loading/collating etc. With this, the TPU can be busy all the time doing the actual computation and won't be blocked by processing and loading.

But I see the problem, so maybe we could use the seed with numpy to make data shuffling reproducible.

cc @patrickvonplaten .

Interesting, what about `run_mlm_flax.py`? I was having a look prior to this PR, and it seems `jax.random.permutation` is also used for data shuffling, or am I missing something?

https://github.com/huggingface/transformers/blob/76c4d8bf26de3e4ab23b8afeed68479c2bbd9cbd/examples/flax/language-modeling/run_mlm_flax.py#L624-L626

That should also be changed. All flax examples are still not completely consistent with each other, which needs to be fixed.

@patil-suraj @patrickvonplaten Is this something you want to change or should I close this PR?

Hey @bminixhofer - sorry for being so slow here. @patil-suraj I'm happy to merge the PR. Think it's good to have 100% reproduciblity with JAX's random seed. I don't think this slows down the script as it's called just once per epoch. If you're ok with the changes feel free to merge 

@patil-suraj - can you take a look here and leave your opinion so that we can resolve the PR? :-)

ping @patil-suraj again

ping @patil-suraj again 

Thanks for merging this! I vaguely remember having problems with batch order with the code as it was previously, but I am not completely sure (it's been some time :sweat_smile: ).