TFWav2Vec2ForCTC: Error when using padded batch and attention mask
Author: yossing-audatic
Pull request: no

## Environment info

- `transformers` version: 4.8.0.dev0
- Platform: Linux
- Python version: 3.8
- PyTorch version (GPU?):  1.8.1
- Tensorflow version (GPU?):  2.4.1
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: Yes 

### Who can help
@patrickvonplaten
@will-rice

Models:

- Wav2Vec2

 -->

## Information

Model I am using: TFWav2Vec2ForCTC 

The problem arises when using:

* Official example script of TFWav2Vec2ForCTC modified to use padded batch

## To reproduce

Steps to reproduce the behavior:

1. Install relevant libraries. 
2. Run code snippet below

```python 
import tensorflow as tf
from transformers import Wav2Vec2Processor, TFWav2Vec2ForCTC
from datasets import load_dataset
import soundfile as sf

processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
model = TFWav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")

def map_to_array(batch):
    speech, _ = sf.read(batch["file"])
    batch["speech"] = speech
    return batch

ds = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
ds = ds.map(map_to_array)

# Pad the speech file with zeros and create corresponding attention mask
speech_len = len(ds["speech"][0])
padded_speech =  ds["speech"][0] + [0.0]*1000
attention_mask = tf.sequence_mask((speech_len,), maxlen=len(padded_speech), dtype=tf.float32)

input_values = processor(padded_speech, return_tensors="tf").input_values # Batch size 1
logits = model(input_values).logits
predicted_ids = tf.argmax(logits, axis=-1)

transcription = processor.decode(predicted_ids[0])

# compute loss
target_transcription = "A MAN SAID TO THE UNIVERSE SIR I EXIST"

# wrap processor as target processor to encode labels
with processor.as_target_processor():
    labels = processor(transcription, return_tensors="tf").input_ids
    

loss = model(input_values, attention_mask=attention_mask, labels=labels).loss
```

## Outputs

```
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-1-263e70f23fd3> in <module>
     33 
     34 
---> 35 loss = model(input_values, attention_mask=attention_mask, labels=labels).loss

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1010         with autocast_variable.enable_auto_cast_variables(
   1011             self._compute_dtype_object):
-> 1012           outputs = call_fn(inputs, *args, **kwargs)
   1013 
   1014         if self._activity_regularizer:

~/code/transformers/src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py in call(self, input_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions, labels, output_hidden_states, return_dict, training, **kwargs)
   1553         )
   1554 
-> 1555         outputs = self.wav2vec2(
   1556             input_values=inputs["input_values"],
   1557             attention_mask=inputs["attention_mask"],

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1010         with autocast_variable.enable_auto_cast_variables(
   1011             self._compute_dtype_object):
-> 1012           outputs = call_fn(inputs, *args, **kwargs)
   1013 
   1014         if self._activity_regularizer:

~/code/transformers/src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py in call(self, input_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict, training, **kwargs)
   1225         hidden_states = self._mask_hidden_states(hidden_states)
   1226 
-> 1227         encoder_outputs = self.encoder(
   1228             hidden_states,
   1229             attention_mask=attention_mask,

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
   1010         with autocast_variable.enable_auto_cast_variables(
   1011             self._compute_dtype_object):
-> 1012           outputs = call_fn(inputs, *args, **kwargs)
   1013 
   1014         if self._activity_regularizer:

~/code/transformers/src/transformers/models/wav2vec2/modeling_tf_wav2vec2.py in call(self, hidden_states, attention_mask, output_attentions, output_hidden_states, return_dict, training)
    993 
    994         if attention_mask is not None:
--> 995             hidden_states = hidden_states * tf.expand_dims(attention_mask, -1)
    996             attention_mask = _expand_mask(attention_mask)
    997         else:

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py in binary_op_wrapper(x, y)
   1162     with ops.name_scope(None, op_name, [x, y]) as name:
   1163       try:
-> 1164         return func(x, y, name=name)
   1165       except (TypeError, ValueError) as e:
   1166         # Even if dispatching the op failed, the RHS may be a tensor aware

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py in _mul_dispatch(x, y, name)
   1494     return sparse_tensor.SparseTensor(y.indices, new_vals, y.dense_shape)
   1495   else:
-> 1496     return multiply(x, y, name=name)
   1497 
   1498 

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/util/dispatch.py in wrapper(*args, **kwargs)
    199     """Call target, and fall back on dispatchers if there is a TypeError."""
    200     try:
--> 201       return target(*args, **kwargs)
    202     except (TypeError, ValueError):
    203       # Note: convert_to_eager_tensor currently raises a ValueError, not a

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/ops/math_ops.py in multiply(x, y, name)
    516   """
    517 
--> 518   return gen_math_ops.mul(x, y, name)
    519 
    520 

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/ops/gen_math_ops.py in mul(x, y, name)
   6066       return _result
   6067     except _core._NotOkStatusException as e:
-> 6068       _ops.raise_from_not_ok_status(e, name)
   6069     except _core._FallbackException:
   6070       pass

/opt/audatic/venv/lib/python3.8/site-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6860   message = e.message + (" name: " + name if name is not None else "")
   6861   # pylint: disable=protected-access
-> 6862   six.raise_from(core._status_to_exception(e.code, message), None)
   6863   # pylint: enable=protected-access
   6864 

/opt/audatic/venv/lib/python3.8/site-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Incompatible shapes: [1,547,768] vs. [1,544,1] [Op:Mul]
```

## Expected behavior
Code should run without errors and produce a loss equivalent to not using padded batch. Without the padded batch, the loss is: 
```python
print(loss)
>>> <tf.Tensor: shape=(), dtype=float32, numpy=39.32432>
```
When using the padded batch and not specifying an attention_mask, the loss is: 
```python 
print(loss)
>>> <tf.Tensor: shape=(), dtype=float32, numpy=39.96655>
```

## Fix:
The bugfix should be quite easy. In line 1217 of transformers.models.wav2vec2.modeling_tf_wav2vec2.py, instead of:
```python 
attention_mask = tf.sequence_mask(output_lengths, dtype=hidden_states.dtype)
```
It should be: 
```python 
max_output_length = self._get_feat_extract_output_lengths(inputs["input_values"].shape[-1])
attention_mask = tf.sequence_mask(output_lengths, maxlen=max_output_length, dtype=hidden_states.dtype)
```




Comments:
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.

Think we can close this one