Flax Support NLLB (or M2M100) model
### Feature request

Add Flax/Jax support for M2M100 so it can be optimize in TPU

### Motivation

NLLB is great model translation that support many languages and also have good accuracy 

it can be used to translate available english large dataset to another language
but this require much resource such as multiple gpu parallelism to cut time for translation

since tpu access more easy to get(through trc program) than multi-gpu access,it would be nice if we have NLLB/M2M100 for flax

to give another reason
flaxmarian in tpu is amazing... it can translate 100k English texts to Spanish in less than 4 minutes using Flax on TPUs with jax parallelism [flax community slack](https://huggingface.slack.com/archives/C025LJDP962/p1626488974341000)

me myself already translated almost 100m sentence using marian flax with 3 days ~~


### Your contribution

although i am not expert yet at flax/jax ,i can make attemp to implement

but i need some pointer how to do that

- AFAIK,NLLB or M2M100 has similar architecture as MBart,since mbart flax has already implemented maybe it can start from that(?)
- need to figure out how to "convert" NLLLB/M2M100 position embedding,etc to jax CMIIW
cc @patil-suraj @sanchit-gandhi in case you'd like to guide @acul3 to contribute the Flax version of M2M100<|||||>Hey @acul3! Awesome, let's do it 💪 Happy to help you through the implementation! On a high-level, the process for adding this model will look something as follows:

1. Copy across the modelling code from Flax Bart
2. Modify the Flax modelling code to match the PyTorch NLLB/M2M100 implementation
\+ write any necessary tests along the way
3. Check whether the Flax logits match with the PyTorch ones
4. Iterate on step 2 until the check in step 3 passes!

Once we have a Flax model that matches the PyTorch logits, we can be confident our implementation is correct :)

As a starting point, we can copy across the Flax Bart modelling code. You can do this through the 'add-new-model-like' command:
https://github.com/huggingface/transformers/tree/main/templates/adding_a_new_model#add-new-model-like-command from the Flax Bart model:
https://github.com/huggingface/transformers/tree/main/src/transformers/models/bart/modeling_flax_bart.py

Once you've done that, feel free to open a WIP PR. We can go from there!<|||||>Hi @sanchit-gandhi 
Thank you for the help

I'll start to work on this today by following the step and open WIP PR first
Will ask some question/pointer after that..thanks again<|||||>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.