Repository: huggingface/transformers (Python)
Task: Bug Fix

Issue:
[SpeechT5] Attention mask not changed according to decoder inputs
### System Info

- `transformers` version: 4.33.3
- Platform: Linux-5.15.0-84-generic-x86_64-with-glibc2.10
- Python version: 3.8.8
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.3.3
- Accelerate version: not installed
- Accelerate config: not found
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>


### Who can help?

@sanchit-gandhi

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

The decoder inputs are changed to be shifted right by one, and interleaved by the reduction factor.
However, the attention mask to the decoder remains the same, which if we use a reduction_factor != 1 will result in a shape missmatch.

You can check the line I am referring to here:
https://github.com/huggingface/transformers/blob/2f3ea08a077ba3133fa8a604b22436cad250b055/src/transformers/models/speecht5/modeling_speecht5.py#L2733

### Expected behavior

The attention mask should have the same changes applied as the decoder input, resulting in the same shape, I believe.


Discussion:
cc @ylacombe could you take a look when you get the chance? You know SpeechT5 pretty well by now!
Hey, thanks for opening this issue!
I will take a look in the next few days, in the meantime, do you have a script to reproduce the mismatch @Joao-Maria-Janeiro ?
Hey @Joao-Maria-Janeiro , any update on a reproducing script?
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/main/CONTRIBUTING.md) are likely to be ignored.
@ylacombe  - +1 This is still an issue. It's very easy to reproduce:

```python
from transformers import SpeechT5Processor, SpeechT5ForSpeechToSpeech
import numpy as np

processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_vc")
model = SpeechT5ForSpeechToSpeech.from_pretrained("microsoft/speecht5_vc")

features = processor(
    audio=[np.random.random(size=(2048,)) for waveform in range(3)],
    audio_target=[np.random.random(size=(2048,)) for waveform in range(3)], 
    return_tensors="pt",
    padding=True,
    sampling_rate=16000,
)
outputs = model(**features, return_dict=True)
```

Produces:

```
Traceback (most recent call last):
  File "[REDACTED]/reproduce.py", line 8, in <module>
    outputs = model(**features, return_dict=True)
  File "/[REDACTED]/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/[REDACTED]/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/[REDACTED]/transformers/models/speecht5/modeling_speecht5.py", line 2953, in forward
    outputs = self.speecht5(
  File "/[REDACTED]/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/[REDACTED]/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/[REDACTED]/transformers/models/speecht5/modeling_speecht5.py", line 2211, in forward
    decoder_outputs = self.decoder(
  File "/[REDACTED]/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/[REDACTED]/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/[REDACTED]/transformers/models/speecht5/modeling_speecht5.py", line 1734, in forward
    outputs = self.wrapped_decoder(
  File "/[REDACTED]/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/[REDACTED]/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/[REDACTED]/transformers/models/speecht5/modeling_speecht5.py", line 1594, in forward
    attention_mask = _prepare_4d_causal_attention_mask(
  File "/[REDACTED]/transformers/modeling_attn_mask_utils.py", line 195, in _prepare_4d_causal_attention_mask
    attention_mask = attn_mask_converter.to_4d(
  File "/[REDACTED]/transformers/modeling_attn_mask_utils.py", line 117, in to_4d
    expanded_4d_mask = expanded_attn_mask if causal_4d_mask is None else expanded_attn_mask + causal_4d_mask
RuntimeError: The size of tensor a (9) must match the size of tensor b (4) at non-singleton dimension 3
```

Patch:
diff --git a/src/transformers/models/speecht5/modeling_speecht5.py b/src/transformers/models/speecht5/modeling_speecht5.py
--- a/src/transformers/models/speecht5/modeling_speecht5.py
+++ b/src/transformers/models/speecht5/modeling_speecht5.py
@@ -64,13 +64,17 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start
     return shifted_input_ids
 
 
-def shift_spectrograms_right(input_values: torch.Tensor, reduction_factor: int = 1):
+def shift_spectrograms_right(
+    input_values: torch.Tensor, reduction_factor: int = 1, attention_mask: Optional[torch.Tensor] = None
+):
     """
     Shift input spectrograms one timestep to the right. Also applies the reduction factor to the sequence length.
     """
     # thin out frames for reduction factor
     if reduction_factor > 1:
         input_values = input_values[:, reduction_factor - 1 :: reduction_factor]
+        if attention_mask is not None:
+            attention_mask = attention_mask[:, reduction_factor - 1 :: reduction_factor]
 
     shifted_input_values = input_values.new_zeros(input_values.shape)
     shifted_input_values[:, 1:] = input_values[:, :-1].clone()
@@ -78,7 +82,7 @@ def shift_spectrograms_right(input_values: torch.Tensor, reduction_factor: int =
     # replace possible -100 values in labels by zeros
     shifted_input_values.masked_fill_(shifted_input_values == -100.0, 0.0)
 
-    return shifted_input_values
+    return shifted_input_values, attention_mask
 
 
 # Copied from transformers.models.wav2vec2.modeling_wav2vec2._compute_mask_indices
@@ -2699,7 +2703,9 @@ def forward(
 
         if labels is not None:
             if decoder_input_values is None:
-                decoder_input_values = shift_spectrograms_right(labels, self.config.reduction_factor)
+                decoder_input_values, decoder_attention_mask = shift_spectrograms_right(
+                    labels, self.config.reduction_factor, decoder_attention_mask
+                )
             if self.config.use_guided_attention_loss:
                 output_attentions = True
 
@@ -3044,7 +3050,9 @@ def forward(
 
         if labels is not None:
             if decoder_input_values is None:
-                decoder_input_values = shift_spectrograms_right(labels, self.config.reduction_factor)
+                decoder_input_values, decoder_attention_mask = shift_spectrograms_right(
+                    labels, self.config.reduction_factor, decoder_attention_mask
+                )
 
         outputs = self.speecht5(
             input_values=input_values,


Test patch:
diff --git a/tests/models/speecht5/test_modeling_speecht5.py b/tests/models/speecht5/test_modeling_speecht5.py
--- a/tests/models/speecht5/test_modeling_speecht5.py
+++ b/tests/models/speecht5/test_modeling_speecht5.py
@@ -909,6 +909,23 @@ def test_model_forward(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model_forward(*config_and_inputs)
 
+    def test_model_forward_with_labels(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs()
+        model = SpeechT5ForTextToSpeech(config=config).to(torch_device).eval()
+
+        input_ids = inputs_dict["input_ids"]
+        attention_mask = inputs_dict["attention_mask"]
+        decoder_attention_mask = inputs_dict["decoder_attention_mask"]
+        labels = inputs_dict["decoder_input_values"]
+
+        result = model(
+            input_ids, attention_mask=attention_mask, labels=labels, decoder_attention_mask=decoder_attention_mask
+        )
+        self.assertEqual(
+            result.spectrogram.shape,
+            (self.model_tester.batch_size, self.model_tester.decoder_seq_length, self.model_tester.num_mel_bins),
+        )
+
     # skipped because there is always dropout in SpeechT5SpeechDecoderPrenet
     def test_decoder_model_past_with_large_inputs(self):
         pass
@@ -1436,6 +1453,23 @@ def test_model_forward(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model_forward(*config_and_inputs)
 
+    def test_model_forward_with_labels(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs()
+        model = SpeechT5ForSpeechToSpeech(config=config).to(torch_device).eval()
+
+        input_values = inputs_dict["input_values"]
+        attention_mask = inputs_dict["attention_mask"]
+        decoder_attention_mask = inputs_dict["decoder_attention_mask"]
+        labels = inputs_dict["decoder_input_values"]
+
+        result = model(
+            input_values, attention_mask=attention_mask, labels=labels, decoder_attention_mask=decoder_attention_mask
+        )
+        self.assertEqual(
+            result.spectrogram.shape,
+            (self.model_tester.batch_size, self.model_tester.decoder_seq_length, self.model_tester.num_mel_bins),
+        )
+
     # skipped because there is always dropout in SpeechT5SpeechDecoderPrenet
     def test_decoder_model_past_with_large_inputs(self):
         pass