FlaxGPT2
# What does this PR do?

[From Patrick:]

This PR adds GPT2 for Flax and lays the groundwork for `generation` in Flax. 
The model is a straightforward translation from PyTorch with two exceptions:

1.) the `past_key_values/caching` is more complex since `generation` has to compile to be usable on TPU with Flax/Jax. This means that the cache cannot grow dynamically, but has to be a static, padded tensor that is filled during generation. This means that the cache can now only be used as follows:
i) decide on a `max_length` to which the model should generate at most and call `model.init_cache(batch_size, max_length)`. This will then add a `"cache"` variable to the expected variables by the model and returns a 0 padded tensor of shape `(batch_size, max_length)`. 
ii) This cache is then subsequently passed & retrieved similar to how it's done in PyTorch:  `logits, past_key_values = model(input_ids, past_key_values=cache)`. When doing so, internally, the `cache` is marked as `mutable` so that the Flax model can step-by-step fill the padded tensor. We also need to pass a `_use_cache` boolean internally to tell the model to initialize the cache in the first place.
iii) The implemented caching system differs slightly from the original flax caching system [here](https://github.com/google/flax/blob/6111a4b2c5cb56520549c6dfc911135fbd8dbb7f/flax/linen/attention.py#L252) since we want to allow the very first decoding pass to take an `input_prompt` that is longer than 1 (*e.g.* for GPT3-like few-shot prompting) - see: https://github.com/google/flax/issues/1317.
iv) A test verifies that compiled generation works as expected
v) This also means that when using the cache, the `attention_mask` has to have the same static shape as the cache, I guess we could have better error messages here to, but I think this should mostly be handled by the generation method which will be added in a follow-up PR: https://github.com/huggingface/transformers/pull/11685.

2.) Since in Flax we rely on `flax.linen.dot_product_attention`, we are adding the attention mask once whereas in PyTorch we are setting certain values to the padding value via `torch.where(...)` [here](https://github.com/huggingface/transformers/blob/c73e35323d5805bff27c5dbd9a5691008be1316a/src/transformers/models/gpt2/modeling_gpt2.py#L187) and are adding large negative numbers [here](https://github.com/huggingface/transformers/blob/c73e35323d5805bff27c5dbd9a5691008be1316a/src/transformers/models/gpt2/modeling_gpt2.py#L191). This is unnecessarily costly in Flax and only matters for unrealistic cases where tokens in the middle of the sequence are padded. For all realistic cases PT and Flax behaves exactly the same, which is shown in the overwritten comparison tests.
@LysandreJik @sgugger - this PR is blocking a couple of other things. Let me know if you guys aren't really ok with some of those things