[BigBird] random attention in FlaxBigBird is not random
### System Info

```shell
NOT NEEDED
```


### Who can help?

@patrickvonplaten 

### Information

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

### Tasks

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

### Reproduction

`numpy.random` is getting used everywhere in `FlaxBigBird` for fetching random indices for computing attention over random tokens. This is wrong as these indices would be cached during jit compilation. Hence, bigbird would attend similar random tokens (i.e same indices) in every step
This can be fixed if random indices are prepared in datacollator and passed to `model.__call__` during training.

https://github.com/huggingface/transformers/blob/ee393c009a243bbb86fa11d2efa771a1704d85cf/src/transformers/models/big_bird/modeling_flax_big_bird.py#L975


This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.<|||||>I see, thanks for putting up this issue here! 

Think the easiest fix in this case would be to use transition from `numpy.random` to `jax.random`: https://jax.readthedocs.io/en/latest/jax.random.html no? This way the user can always pass a PRNG key to the forward which would define the random attention mask?

What do you think @vasudevgupta7 ?<|||||>Hello @patrickvonplaten,

sorry for the late reply. I missed your comment somehow.

> Think the easiest fix in this case would be to use transition from numpy.random to jax.random: https://jax.readthedocs.io/en/latest/jax.random.html no? This way the user can always pass a PRNG key to the forward which would define the random attention mask?

Yes, this should be the best way to go. But then Jax implementation will diverge from PyTorch implementation in case of block sparse attention right?? and some tests would fail ig.<|||||>Thanks for the reply @vasudevgupta7!

Leaving this as a "second good issue" now as I won't find the time to dig into it in the near future. @community Please ping me here if you'd like this to be fixed soon<|||||>Hi @patrickvonplaten @thevasudevgupta - is this still a problem? If so I would love to pick it up to have some work on the Christmas break :)<|||||>This would be great @Bearnardd :-) <|||||>@thevasudevgupta this issue can be closed as the fix is merged :)<|||||>Very cool to know that!