Flax Big Bird
# What does this PR do?

--------------------------------------------------------------------
**🚨 Bug detection 🚨**
`BigBirdForMultipleChoice` was incorrect and is corrected in this PR. This is a breaking change for all BigBird models that have been trained on multiple choice (0 on the hub currently)
--------------------------------------------------------------------

This PR will add `FlaxBigBirdModel`.

Evaluation Notebook: https://colab.research.google.com/drive/1rx_G9awurQekrK1mSzd3A9F_UciTTuTY?usp=sharing#scrollTo=ecjNtnAuKYo8

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

## Before submitting
- [x] 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?
- [x] 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.
- [x] 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).
- [x] Did you write any new necessary tests?

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**. Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

@patrickvonplaten 
Wow awesome work! I think the important next step would be to make the test:

```
tests/test_modeling_flax_bigbird.py::FlaxBigBirdModelTest::test_jit_compilation
```

If this test works, we can enable super fast training on TPU :-) <|||||>This test is passing for all model classes except for `FlaxBigBirdForMultipleChoice` and models are already `jit` compatible.  It's failing since there is some bug in `FlaxBertForMultipleChoice` (BERT) which doesn't allow it to work for `seqlen > 1`.<|||||>> This test is passing for all model classes except for `FlaxBigBirdForMultipleChoice` and models are already `jit` compatible. It's failing since there is some bug in `FlaxBertForMultipleChoice` (BERT) which doesn't allow it to work for `seqlen > 1`.

I'm sorry - I don't follow here 100%. Running `RUN_PT_FLAX_CROSS_TESTS=1 pytest tests/test_modeling_flax_bert.py` passes, so `FlaxBertForMultipleChoice` seems to work correctly. Can you maybe open an issue showing the problem with `FlaxBertForMultipleChoice`? <|||||>Hey Vasu, 

could you remove the `.ipynb` debugger file? :-)<|||||>done.<|||||>@LysandreJik @sgugger @stas00 

The jitted-Flax tests are getting too expensive to be run at every commit:
```
423.98s call     tests/test_modeling_flax_big_bird.py::FlaxBigBirdModelTest::test_jit_compilation
64.09s call     tests/test_modeling_flax_bert.py::FlaxBertModelTest::test_jit_compilation
51.94s call     tests/test_modeling_flax_electra.py::FlaxElectraModelTest::test_jit_compilation
43.62s call     tests/test_modeling_flax_roberta.py::FlaxRobertaModelTest::test_jit_compilation
37.21s call     tests/test_modeling_flax_big_bird.py::FlaxBigBirdModelTest::test_hidden_states_output
29.00s call     tests/test_modeling_flax_gpt2.py::FlaxGPT2ModelTest::test_greedy_generate
28.08s call     tests/test_modeling_flax_big_bird.py::FlaxBigBirdModelTest::test_model_outputs_equivalence
27.02s call     tests/test_modeling_flax_gpt2.py::FlaxGPT2ModelTest::test_sample_generate_logits_warper
25.48s call     tests/test_modeling_flax_bert.py::FlaxBertModelTest::test_attention_outputs
25.26s call     tests/test_modeling_flax_gpt2.py::FlaxGPT2ModelTest::test_sample_generate
25.10s call     tests/test_modeling_flax_gpt2.py::FlaxGPT2ModelTest::test_greedy_generate_attn_mask
24.77s call     tests/test_modeling_flax_gpt2.py::FlaxGPT2ModelTest::test_sample_generate_attn_mask
23.91s call     tests/test_modeling_flax_electra.py::FlaxElectraModelTest::test_attention_outputs
23.50s call     tests/test_modeling_flax_clip.py::FlaxCLIPModelTest::test_jit_compilation
20.20s call     tests/test_modeling_flax_roberta.py::FlaxRobertaModelTest::test_attention_outputs
16.70s call     tests/test_modeling_flax_bert.py::FlaxBertModelTest::test_model_outputs_equivalence
15.46s call     tests/test_modeling_flax_gpt2.py::FlaxGPT2ModelTest::test_jit_compilation
13.64s call     tests/test_modeling_flax_clip.py::FlaxCLIPVisionModelTest::test_jit_compilation
12.68s call     tests/test_modeling_flax_bert.py::FlaxBertModelTest::test_hidden_states_output
12.43s call     tests/test_modeling_flax_electra.py::FlaxElectraModelTest::test_model_outputs_equivalence
12.05s call     tests/test_tokenization_mbart50.py::MBartTokenizationTest::test_save_pretrained
11.96s call     tests/test_modeling_flax_big_bird.py::FlaxBigBirdModelTest::test_forward_signature
11.88s call     tests/test_modeling_flax_roberta.py::FlaxRobertaModelTest::test_model_outputs_equivalence
11.63s call     tests/test_modeling_flax_clip.py::FlaxCLIPVisionModelTest::test_attention_outputs
11.11s call     tests/test_modeling_flax_clip.py::FlaxCLIPModelTest::test_get_image_features
```

However they are super important to ensure that the model works on TPU. Can we somehow run them only on approval or it's probably easier to just set them to "slow" for now?<|||||>@vasudevgupta7 could you also run `make style` one last time? :-)<|||||>> The jitted-Flax tests are getting too expensive to be run at every commit:
[...]
> 
> However they are super important to ensure that the model works on TPU. Can we somehow run them only on approval or it's probably easier to just set them to "slow" for now?

What you're saying is that the TPU tests won't be then run at all, because we don't have a TPU runner for slow tests, correct?

I think Circle CI has a mechanism where you can trigger certain runs by adding a special keyword to the commit message.  But that might be too complicated to remember to do.

How about this idea. Leave an open PR with a circle-ci jitted-flax tests job that only exists in this PR, and which gets rebased on a nightly basis through a cron-job and pushed, that would give a poor man's scheduled CI run on TPU. Perhaps there are some easier ways.<|||||>We can mark them as slow if you want but there's no Flax GPU CI right now, so the slow tests won't be run :) @stas00's proposal sounds good!<|||||>tests failing on CircleCI are unrelated to this PR.<|||||>Flax tests are disabled now - opening a PR that will run them as proposed by @stas00 <|||||>Merging - great job @vasudevgupta7 