Flax Speech-Encoder-Decoder Model
Author: sanchit-gandhi
Pull request: yes

# What does this PR do?

Currently, Speech-Encoder-Decoder Models are only available in PyTorch. This PR adds the **Flax** Speech-Encoder-Decoder Model and all accompanying tests.

The accuracy of the Flax Speech-Encoder-Decoder Model is validated through cross-test comparison with the corresponding PyTorch Speech-Encoder-Decoder Model. The final logit values agree to within a `4e-2` threshold, making certain that the Flax Model operates as intended.

Comments:
The issue being encountered occurs when the model weights are initialized: the `encoder_hidden_states` don’t match the shape of the `decoder_hidden_states`, meaning that the cross-attention layers throw an error when the decoder is initialised. I have been inspecting the code to try and establish where the two deviate in shape. For reference, running this code produces the error:

```py
from transformers import FlaxSpeechEncoderDecoderModel
# Initialize a wav2vec2-2-gpt2 from pretrained wav2vec2 and gpt2 models
# Note that the cross-attention layers will be randomly initialized
model_fx = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained("patrickvonplaten/dummy_wav2vec2_with_adapter", "gpt2", encoder_from_pt=True)
```

> The issue being encountered occurs when the model weights are initialized: the `encoder_hidden_states` don’t match the shape of the `decoder_hidden_states`, meaning that the cross-attention layers throw an error when the decoder is initialised. I have been inspecting the code to try and establish where the two deviate in shape. For reference, running this code produces the error:
> 
> ```python
> from transformers import FlaxSpeechEncoderDecoderModel
> # Initialize a wav2vec2-2-gpt2 from pretrained wav2vec2 and gpt2 models
> # Note that the cross-attention layers will be randomly initialized
> model_fx = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained("patrickvonplaten/dummy_wav2vec2_with_adapter", "gpt2", encoder_from_pt=True)
> ```

This already looks like a pretty edge-case loading. Could we maybe in a first step try to see whether the following works:

```py
from transformers import FlaxSpeechEncoderDecoderModel
# Initialize a wav2vec2-2-gpt2 from pretrained wav2vec2 and gpt2 models
# Note that the cross-attention layers will be randomly initialized
model_fx = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained("patrickvonplaten/dummy_wav2vec2_with_adapter", "gpt2")
```
? 

Meaning that the weights don't have to be converted from PT to Flax. For this you probably first need to do:

```py
from transformers import FlaxWav2Vec2Model

model = FlaxWav2Vec2Model.from_pretrained("patrickvonplaten/dummy_wav2vec2_with_adapter", from_pt=True)
model.push_to_hub("<your-url>")
```

Thanks for the reply! Converting the PT weights to Flax and re-running the script still throws the error:
```py
from transformers import FlaxSpeechEncoderDecoderModel
# Initialize a wav2vec2-2-gpt2 from pretrained **Flax** wav2vec2 and gpt2 models
# Note that the cross-attention layers will be randomly initialized
model_fx = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained("sanchit-gandhi/dummy_flax_wav2vec2_with_adapter", "gpt2")
```
The first forward pass in initialising the model is performed successfully: the shape of both the `encoder_hidden_states` and `decoder_hidden_states` is `(1, 1, 768)`. In the second forward pass, the `encoder_hidden_states` are `(1, 2, 768)`, and the `decoder_hidden_states` are `(1, 1024, 768)`. Not sure where this discrepancy in `seq_len` has come about. They should both be `(1, 1, 768)` I believe.

@sanchit-gandhi,

Just pushed a fix for the problem. The reason was that the encoder downsamples the seq_len while the decoder does not. This means that we can't have the same encoder and decoder input lengths. 

The script still fails, but now because the attention_mask is not correct. Could you take a look here? We should implement `_get_feature_vector_attention_mask` now as you correctly noted :-) 

Quick question about our PT-Flax cross-test thresholds. I noticed that in the `test_modeling_flax_common.py`, we use a threshold of 4e-2 when comparing Flax and PT outputs: https://github.com/huggingface/transformers/blob/60ba48205e4ec070780e8cb8d461421b77432bad/tests/test_modeling_flax_common.py#L192 whereas in `test_modeling_flax_encoder_decoder.py` the threshold is set much more aggressively at 1e-5: https://github.com/huggingface/transformers/blob/60ba48205e4ec070780e8cb8d461421b77432bad/tests/test_modeling_flax_encoder_decoder.py#L261 Is there a reason why the threshold is set so much lower in the latter?

> that

Good question! The reason is mostly due to different frameworks and how the matrix multiplication is done. In JAX a LAX-backend implementation of `matmul()` is used which approximates the output matrix often instead of doing the actual computation - see here: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.matmul.html

So if we see differences of like 0.05 we are still pretty sure that the model is correctly implemented in Flax