[Flax] Change all `shift_tokens_right` to numpy code
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.