Fix TFWav2Vec2 SpecAugment
Author: will-rice
Pull request: yes

# What does this PR do?
This PR fixes the SpecAugment implementation for TFWav2Vec2. I had a lot of trouble with this during the original PR, so I'm not 100% if this is correct. I would love feedback because at this point there must be a knowledge gap. 

Fixes # (issue)
https://github.com/huggingface/transformers/issues/12264

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@patrickvonplaten 


Comments:
Based on the original paper the time mask should be applied on the time axis like 
![image](https://user-images.githubusercontent.com/25072137/122934746-cebeab00-d33d-11eb-92bd-c18652a107e9.png)
However, wav2vec2 is masking hidden states. In principle, it should work the same way and if correctly applied you should be able to see the zeroed span in the time dimension of the hidden states.

Hey @will-rice,

Thanks for fixing the bug. As I understand it, when applying spec_augment along the time axis we are actually not setting the values to zero but to a trained mask embedding vector (this is mostly because that's how it's used for pretraining actually). 
IMO, to correctly implement this one has to make use of `tf.where`. A good way would be to expand the masked indices and `self.masked_spec_embed` to be of the same size as `hidden_states` (as explained above) ati use `tf.where`.

For spec_augment along the feature axis I would also suggest to use `tf.where` -> expand the feature indices (this time along the time axis (seq length) and then one can simply do:

```
hidden_states = tf.where(expanded_mask, hidden_states, 0)
```

Let me know if this is not understandable ;-) 

> Hey @will-rice,
> 
> Thanks for fixing the bug. As I understand it, when applying spec_augment along the time axis we are actually not setting the values to zero but to a trained mask embedding vector (this is mostly because that's how it's used for pretraining actually).
> IMO, to correctly implement this one has to make use of `tf.where`. A good way would be to expand the masked indices and `self.masked_spec_embed` to be of the same size as `hidden_states` (as explained above) ati use `tf.where`.
> 
> For spec_augment along the feature axis I would also suggest to use `tf.where` -> expand the feature indices (this time along the time axis (seq length) and then one can simply do:
> 
> ```
> hidden_states = tf.where(expanded_mask, hidden_states, 0)
> ```

"fixing" :stuck_out_tongue:. I believe I understand the way to do it now but will post questions here if I get stuck. Thank you for walking through this!