Add Adapter Weighs to Flax
Author: patrickvonplaten
Pull request: no

# 🚀 Feature request

Currently it's possible to add an adapter on the top of PyTorch Wav2Vec2 (https://github.com/huggingface/transformers/blob/1d94d575461a76cb1dcb3ebe6e85f1c85d1dafcd/src/transformers/models/wav2vec2/modeling_wav2vec2.py#L1033) - however an equivalent module is missing in Flax: https://github.com/huggingface/transformers/blob/master/src/transformers/models/wav2vec2/modeling_flax_wav2vec2.py. 

The adapter is essentially used to reduce the time dimension further so that the encoder's output hidden states have a time context window which is more similar to that of a subword token instead of just a character (as done for CTC). This was introduced for the XLS-R paper: https://arxiv.org/abs/2111.09296 and can be found in the original fairseq code here: https://github.com/pytorch/fairseq/blob/5d2be954bb7531bff92c195e61aa50a8ddd0baab/fairseq/models/speech_to_text/xm_transformer.py#L245

We should add this to Flax as well for the Seq2Seq experiments.

## Goal the following script should give identical results:

```python
import torch
import numpy as np
from transformers import FlaxWav2Vec2Model, Wav2Vec2Model

model_fx = FlaxWav2Vec2Model.from_pretrained("patrickvonplaten/dummy_wav2vec2_with_adapter", from_pt=True)
model_pt = Wav2Vec2Model.from_pretrained("patrickvonplaten/dummy_wav2vec2_with_adapter")

input_torch = torch.ones((2, 5000), dtype=torch.float32)
input_fx = input_torch.cpu().numpy()

with torch.no_grad():
    output_logits_pt = model_pt(input_torch).last_hidden_state
output_logits_flax = model_fx(input_fx).last_hidden_state

print("Check if shapes are equal")
print(f"Shape PyTorch {output_logits_pt.shape} | Shape Flax {output_logits_flax.shape}")

print("Check if output values are equal")
print(f"Diff {np.max(np.abs(output_logits_pt.numpy()) - np.asarray(np.abs(output_logits_flax)))})")
```

This script fails at the moment because both the shape and the output logits are different. You can also see when loading the model in Flax that some weights are not used since the implementation of FlaxWav2Vec2Adaptor is missing.

Traceback:
```bash
Some weights of the model checkpoint at patrickvonplaten/dummy_wav2vec2_with_adapter were not used when initializing FlaxWav2Vec2Model: {('adapter', 'layers', '2', 'conv', 'bias'), ('adapter', 'layers', '1', 'conv', 'bias'), ('adapter', 'layers', '0', 'conv', 'kernel'), ('adapter', 'layers', '1', 'conv', 'kernel'), ('adapter', 'layers', '2', 'conv', 'kernel'), ('adapter', 'layers', '0', 'conv', 'bias')}
- This IS expected if you are initializing FlaxWav2Vec2Model from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing FlaxWav2Vec2Model from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Check if shapes are equal
Shape PyTorch torch.Size([2, 2, 768]) | Shape Flax (2, 15, 768)
Check if output values are equal
Traceback (most recent call last):
  File "/home/patrick/./wav2vec2_flax_add_adapter.py", line 20, in <module>
    print(f"Diff {np.max(np.abs(output_logits_pt.numpy()) - np.asarray(np.abs(output_logits_flax)))})")
ValueError: operands could not be broadcast together with shapes (2,2,768) (2,15,768) 
```

Comments:
@sanchit-gandhi