Integrate Bert-like model on Flax runtime.
Author: mfuntowicz
Pull request: yes

This Pull Request attempts to bring support for [Flax](https://github.com/google/flax) framework as part of transformers. 

Main focus as been put on providing BERT-like models, principally by making it possible to load PyTorch checkpoints and doing the necessary conversions (few) directly on the fly. Supports also providing a **msgpack** formatted file from Flax.

`save_pretrained` will save the model through **msgpack** format to avoid dependency on torch inside Jax code. 

**Targeted models:** 
- [x] Bert
- [x] RoBERTa 
- [ ] DistilBERT
- [ ] DistilRoBERTa

**If not too hard**

- [ ] CamemBERT



Comments:
As you said, Jax is a library that interact with numpy to provide additional features: autodiff, auto-vectorization [(**vmap**)](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html#Auto-vectorization-with-vmap) and auto-parallelization [(**pmap**)](https://jax.readthedocs.io/en/latest/jax.html#parallelization-pmap). 

Jax is essentially stateless, which is reflected here through the function to differentiate (the model) doesn't holds the parameters. They have to be referenced somewhere else and feed somehow.

`JaxPreTrainedModel` is introduced here mainly to handle the serialization of such model and provide conversion. Also, one specificity of Jax is many different Neural Network library are currently being implemented on top of it: 

- Google Flax (https://github.com/google/flax)
- Google Trax (https://github.com/google/trax)
- DeepMind Haiku (https://github.com/deepmind/dm-haiku)

In that aspect, @madisonmay is currently working on a [Haiku Bert integration](https://github.com/huggingface/transformers/pull/3520) in transformers. My hope it to be able to share as many things as possible between the two implementations (_but can't be sure for now_) 

Alright, that makes sense. Thanks for the explanation.

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Unstale

# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/3722?src=pr&el=h1) Report
> Merging [#3722](https://codecov.io/gh/huggingface/transformers/pull/3722?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/60de910e6010c76c25dd0ed0999e4c69f9692371?el=desc) will **increase** coverage by `2.55%`.
> The diff coverage is `90.11%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/3722/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/3722?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #3722      +/-   ##
==========================================
+ Coverage   78.32%   80.88%   +2.55%     
==========================================
  Files         187      165      -22     
  Lines       37162    30383    -6779     
==========================================
- Hits        29107    24575    -4532     
+ Misses       8055     5808    -2247     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/3722?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/\_\_init\_\_.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9fX2luaXRfXy5weQ==) | `99.30% <ø> (-0.11%)` | :arrow_down: |
| [src/transformers/modeling\_flax\_auto.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19mbGF4X2F1dG8ucHk=) | `60.86% <60.86%> (ø)` | |
| [src/transformers/modeling\_flax\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19mbGF4X3V0aWxzLnB5) | `83.60% <83.60%> (ø)` | |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `82.92% <92.85%> (-0.05%)` | :arrow_down: |
| [src/transformers/modeling\_flax\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19mbGF4X3JvYmVydGEucHk=) | `94.11% <94.11%> (ø)` | |
| [src/transformers/modeling\_flax\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19mbGF4X2JlcnQucHk=) | `96.50% <96.50%> (ø)` | |
| [src/transformers/testing\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90ZXN0aW5nX3V0aWxzLnB5) | `67.66% <100.00%> (+0.38%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_mobilebert.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9tb2JpbGViZXJ0LnB5) | `24.55% <0.00%> (-72.40%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_flaubert.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9mbGF1YmVydC5weQ==) | `24.53% <0.00%> (-65.14%)` | :arrow_down: |
| [src/transformers/trainer.py](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90cmFpbmVyLnB5) | `55.11% <0.00%> (-9.71%)` | :arrow_down: |
| ... and [156 more](https://codecov.io/gh/huggingface/transformers/pull/3722/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/3722?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/3722?src=pr&el=footer). Last update [60de910...c0d1c81](https://codecov.io/gh/huggingface/transformers/pull/3722?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


cc @levskaya

It looks like a file is missing:

```
$ make fixup
[...]
Checking all models are properly tested.
Traceback (most recent call last):
  File "utils/check_repo.py", line 327, in <module>
    check_repo_quality()
  File "utils/check_repo.py", line 321, in check_repo_quality
    check_all_models_are_tested()
  File "utils/check_repo.py", line 212, in check_all_models_are_tested
    new_failures = check_models_are_tested(module, test_file)
  File "utils/check_repo.py", line 182, in check_models_are_tested
    tested_models = find_tested_models(test_file)
  File "utils/check_repo.py", line 163, in find_tested_models
    with open(os.path.join(PATH_TO_TESTS, test_file)) as f:
FileNotFoundError: [Errno 2] No such file or directory: 'tests/test_modeling_flax_utils.py'
Makefile:25: recipe for target 'extra_quality_checks' failed
make: *** [extra_quality_checks] Error 1
```

Shouldn't the CI have caught this?

Looks like a problem in `make fixup`, `make quality` runs fine (and that's what the CI runs).

Nope, both run the same sub-target: `extra_quality_checks`

```
$ make quality
[...]
python utils/check_copies.py
python utils/check_dummies.py
python utils/check_repo.py
2020-10-19 12:11:26.345843: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
Checking all models are properly tested.
Traceback (most recent call last):
  File "utils/check_repo.py", line 327, in <module>
    check_repo_quality()
  File "utils/check_repo.py", line 321, in check_repo_quality
    check_all_models_are_tested()
  File "utils/check_repo.py", line 212, in check_all_models_are_tested
    new_failures = check_models_are_tested(module, test_file)
  File "utils/check_repo.py", line 182, in check_models_are_tested
    tested_models = find_tested_models(test_file)
  File "utils/check_repo.py", line 163, in find_tested_models
    with open(os.path.join(PATH_TO_TESTS, test_file)) as f:
FileNotFoundError: [Errno 2] No such file or directory: 'tests/test_modeling_flax_utils.py'
Makefile:25: recipe for target 'extra_quality_checks' failed
make: *** [extra_quality_checks] Error 1
```

This is with the latest master.

PR with fix https://github.com/huggingface/transformers/pull/7914

The question is - why CI didn't fail? It reports no problem here:
https://app.circleci.com/pipelines/github/huggingface/transformers/14040/workflows/6cd2b931-ce7e-4e99-b313-4a34326fcece/jobs/101513

Once I got this fixed, 2 more issues came up:

```
python utils/check_repo.py
2020-10-19 12:22:10.636984: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
Checking all models are properly tested.
Traceback (most recent call last):
  File "utils/check_repo.py", line 328, in <module>
    check_repo_quality()
  File "utils/check_repo.py", line 322, in check_repo_quality
    check_all_models_are_tested()
  File "utils/check_repo.py", line 217, in check_all_models_are_tested
    raise Exception(f"There were {len(failures)} failures:\n" + "\n".join(failures))
Exception: There were 2 failures:
test_modeling_flax_bert.py should define `all_model_classes` to apply common tests to the models it tests. If this intentional, add the test filename to `TEST_FILES_WITH_NO_COMMON_TESTS` in the file `utils/check_repo.py`.
test_modeling_flax_roberta.py should define `all_model_classes` to apply common tests to the models it tests. If this intentional, add the test filename to `TEST_FILES_WITH_NO_COMMON_TESTS` in the file `utils/check_repo.py`.
Makefile:25: recipe for target 'extra_quality_checks' failed

```
Fixed in the same PR.