[Whisper] Make tests faster

[sanchit-gandhi]

What does this PR do?

Reduces the input seq length of the Whisper tests from 1500 -> 60 frames. This in turn should speed up the tests quite considerably.

[sanchit-gandhi]

We have to override this method to ensure that we init the Flax weights with the downsampled sequence length correctly (e.g. pass input_shape=init_shape)

[ydshieh]

I am a bit confused here. If you look at FlaxWhisperModelTest, there is no such overriding to pass input_shape. However, FlaxWhisperModelTester uses the low number as in this PR.

Why we don't need to pass init_shape in FlaxWhisperModelTest?

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[ydshieh]

Sorry, I am asking not because I see a test being slow, but just I saw some more Whisper test failures on daily CI, which is tests/models/whisper/test_modeling_whisper.py::WhisperModelTest::test_cpu_offload.

But yes, in general, it's best to use low number. I will take a look.

[sanchit-gandhi]

Note that the Whisper tests have already been flagged as being slow (#23736) so this should help combat this issue!

[ydshieh]

It's not because it's slow test that we use large value without really valid reason :-). Always better to make them use low values is the goal, unless it's absolute necessary.

I still have questions on why we don't need to pass input_shape in the corresponding flax test file.

[ydshieh]

OK, in flax test file, I see

        self.all_model_classes = (
            make_partial_class(model_class, input_shape=self.init_shape) for model_class in self.all_model_classes
        )

probably it's the reason.

[ydshieh]

Thank you @sanchit-gandhi LGTM. I convinced myself regarding the input_shape thing.
But any comment is welcomed.

[sanchit-gandhi]

Yep agreed - the seq len was unnecessarily high here :) You're spot on regarding the init shape: we have to change this based on the sequence length since Flax Whisper initialises the positional embeddings based on the context window, so if we change the seq len (= context window) we need to init the weights with the new shape

[sgugger]

Thanks!