updating _compute_mask_indices fn to work with torch compile
Author: Kirandevraj
Pull request: yes

fixes #22849 

The inplace operations are replaced with out-of-place ones to fix the torch compile computational graph breakage.
This method converts the numpy operations to torch operations in _compute_mask_indices function.
The _compute_mask_indices is used when using SpecAugmentation in  wav2vec2 training.

@sanchit-gandhi 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_24668). All of your documentation changes will be reflected on that endpoint.


> And also add a new test to check that compiling the forward call works when we have spec aug activated

We already have a test case ```test_mask_time_prob_ctc``` that check the forward call with spec aug activated. 
Do you mean when using compile mode - we want to have a test case?

> Do you mean when using compile mode - we want to have a test case?

Yes please - one test to make sure this PR gives the expected behaviour would be grand!

I believe we just need an end-to-end test here and then we're good to go right @Kirandevraj?

Yes, 
Also other test cases that expects ```np.ndarray``` output from the ```_compute_mask_indices``` have to updated to expect the new output ```torch.Tensor``` from ```_compute_mask_indices```

Awesome - this is super close to completion then! Would you like to see it home? Feel free to ping me with any other questions / queries, more than happy to help! 

Yes, I am working on the test case. let me share it soon.

When I use ```sampled_negative_indices``` in test cases. I am facing the following error from the last line.
```AssertionError('Loop-carried variable _tmp1 has initial type <[1, 2048], int1> but is re-assigned to  <[1, 2048], int8> in loop! Please make sure that the type stays consistent.')``` error by the 'inductor'. 
Sample code: 
```
import torch
from transformers import AutoFeatureExtractor, Wav2Vec2ForPreTraining
from transformers.models.wav2vec2.modeling_wav2vec2 import _compute_mask_indices, _sample_negative_indices
from datasets import load_dataset

if torch.cuda.is_available():
    device = torch.device("cuda")
    print("CUDA is available")
else:
    device = torch.device("cpu")
    print("CUDA is not available")


feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base")
model = Wav2Vec2ForPreTraining.from_pretrained("facebook/wav2vec2-base")
model.to(device)

ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
input_values = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt").input_values.to(device)  # Batch size 1

# compute masked indices
batch_size, raw_sequence_length = input_values.shape
sequence_length = model._get_feat_extract_output_lengths(raw_sequence_length).item()
mask_time_indices = _compute_mask_indices(shape=(batch_size, sequence_length), mask_prob=0.2, mask_length=2)

sampled_negative_indices = _sample_negative_indices(
features_shape=(batch_size, sequence_length),
num_negatives=model.config.num_negatives,
mask_time_indices=mask_time_indices,
)

mask_time_indices = torch.tensor(data=mask_time_indices, device=input_values.device, dtype=torch.long)
sampled_negative_indices = torch.tensor(data=sampled_negative_indices, device=input_values.device, dtype=torch.long)

compiled_model = torch.compile(model)
outputs = compiled_model(input_values,mask_time_indices=mask_time_indices, sampled_negative_indices=sampled_negative_indices)
```
I believe this is from ```if sampled_negative_indices is not None:``` sub-block code snippet in forward function from wav2vec2. 
Should we go about fixing it to make it work with torch compile or can we create test case without using the ```sampled_negative_indices```
Can you direct me where the assertion error is getting raised to inquire.

Do you have a full stack trace @Kirandevraj? This should give us a clue as to where the dtype is changing. IMO it could be worth going through the `sampled_negative_indices` function with a debugger and just checking that the input/output dtypes are consistent. The error message suggests we're doing an un-intentional upcast/downcast somewhere in the function

The above error was from ```neg_is_pos.any()```. I updated the code to ```neg_is_pos.sum() > 0``` to work with torch compile.
I have added the test case ```test_torch_compiled_model_for_pretraining``` taking inspiration from ```test_model_for_pretraining```. I have updated few inplace operation with ```torch.where``` function. 

Awesome, sounds great! Let me know when you'd like me to take a second look at the code + tests! Getting close to merge here!

These are the four test errors that I am facing in my brach and the main branch in my system and all the other test cases are passing from the test_modeling_wav2vec2.py
```
FAILED tests/models/wav2vec2/test_modeling_wav2vec2.py::Wav2Vec2ModelTest::test_initialization - AssertionError: 0.5709530115127563 not found in [0.0, 1.0] : Parameter wav2vec2.encoder.pos_conv_embed.conv.parametrizations.weight.original0 of model <class 'transformers.models.wav2vec2.modeling_wav2vec2.Wa...
FAILED tests/models/wav2vec2/test_modeling_wav2vec2.py::Wav2Vec2ModelTest::test_torch_fx - AssertionError: Couldn't trace module: symbolically traced variables cannot be used as inputs to control flow
FAILED tests/models/wav2vec2/test_modeling_wav2vec2.py::Wav2Vec2ModelTest::test_torch_fx_output_loss - AssertionError: Couldn't trace module: symbolically traced variables cannot be used as inputs to control flow
FAILED tests/models/wav2vec2/test_modeling_wav2vec2.py::Wav2Vec2RobustModelTest::test_initialization - AssertionError: 0.585796058177948 not found in [0.0, 1.0] : Parameter wav2vec2.encoder.pos_conv_embed.conv.parametrizations.weight.original0 of model <class 'transformers.models.wav2vec2.modeling_wav2vec2.Wav...
```
Can you take a look.

Hey @Kirandevraj! The first two errors look like model initialisation errors (the dynamic range of the weights isn't in the expected range). These should be quite easy to fix as a last step, but I don't think we've changed anything in the state dict or the initialisation process, so shouldn't expect them to change?

The other two failing tests suggest that the model is not working with torch fx tracing. Could you try running the model with trace enabled?

```python
from transformers import Wav2Vec2Model, Wav2Vec2FeatureExtractor
from transformers.utils.fx import symbolic_trace
import numpy as np

model = Wav2Vec2Model.from_pretrained("hf-internal-testing/tiny-random-wav2vec2")
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2")

raw_speech = np.ones(16000)

inputs = feature_extractor(raw_speech, sampling_rate=16000, return_tensors="pt")
input_names = list(inputs.keys())

traced_model = symbolic_trace(model, input_names)
traced_output = traced_model(inputs)
```

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.