Make Flax pt-flax equivalence test more aggressive
Author: ydshieh
Pull request: yes

# What does this PR do?

Make Flax pt-flax equivalence test more aggressive. (Similar to #15839 for PT/TF).

It uses `output_hidden_states=True` and `output_attentions=True` to test all output tensors (in a recursive way).

Also, it lowers the tolerance from `4e-2` to `1e-5`. (From the experience I gained in PT/TF test, if an error > `1e-5`, I always found a bug to fix).           

(A bit) surprisingly, but very good news: unlike PT/TF, there is no PT/Flax inconsistency found by this more aggressive test! (@patil-suraj  must have done a great job on flax models :-) )

Flax: @patil-suraj @patrickvonplaten 


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_15841). All of your documentation changes will be reflected on that endpoint.

Also very surprised that the tests are all passing. @ydshieh - can you double check that the tests are actually passing for most models. Some edge-case models that could be tested locally:
- BigBird
- Pegagus
- GPT2

@patil-suraj @patrickvonplaten 

There are a few things to fix in this PR (same for the PT/TF) - some tests are just ignored by my mistakes.
I will let you know once I fix them.

**[Updated Info.]**

- The (more aggressive) PT/TF test is merged to `master`
- I fixed some bugs for this new PT/Flax test
  - There are 10 failures
    - 6 have `very large` difference between PT/Flax ( > `1.8`) --> I will check if I can fix them easily
    - 4 have `large` difference (0.01 ~ 0.05) --> Need to verify if these are expected 
 

>         * 6 have `very large` difference between PT/Flax ( > `1.8`) --> I will check if I can fix them easily
>         * 4 have `large` difference (0.01 ~ 0.05) --> Need to verify if these are expected
~~Good news! Once #16167 and #16168 are merged, this more aggressive PT/Flax test will pass with `1e-5` on `CPU`. I will test with `GPU` later.~~

Sorry, but please ignore the above claim. The tests ran flax models on CPU because the GPU version of Jax/Flax were not installed.

----
Running on `GPU` with `1e-5` also passes! (run 3 times per model class)

_The documentation is not available anymore as the PR was closed or merged._

Update:

After rebasing on a more recent commit on master (`5a6b3ccd28a320bcde85190b0853ade385bd4158`), this test with `1e-5` work fine!

I ran this new test on GPU (inside docker container that we use for CI GPU testing + with `jax==0.3.0`). The only errors I got is

- Flax/Jax  `failed to determine best cudnn convolution algorithm for ...`
  - Need to find out the cause eventually.

Think this PR is ready! @patil-suraj @patrickvonplaten 

(After installing `jax[cuda11_cudnn805]` instead of `jax[cuda11_cudnn82]`, the errors listed below no longer appear)

----

## Error logs
```
FAILED tests/beit/test_modeling_flax_beit.py::FlaxBeitModelTest::test_equivalence_flax_to_pt - RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
FAILED tests/clip/test_modeling_flax_clip.py::FlaxCLIPVisionModelTest::test_equivalence_flax_to_pt - RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
FAILED tests/clip/test_modeling_flax_clip.py::FlaxCLIPModelTest::test_equivalence_flax_to_pt - RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
FAILED tests/vit/test_modeling_flax_vit.py::FlaxViTModelTest::test_equivalence_flax_to_pt - RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
FAILED tests/wav2vec2/test_modeling_flax_wav2vec2.py::FlaxWav2Vec2ModelTest::test_equivalence_flax_to_pt - RuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
```

It's all green! I will approve my own PR too: LGTM!