[Flax] Change all shift_tokens_right to numpy code

[patrickvonplaten]

shift labels right is called usually in the data collator and therefore should not be written in jax, but in numpy to not block the TPU. We should make sure that all Encoder-Decoder models have their shift_tokens_right implemented in numpy as it's faster.