Hey @gau-nernst - thanks very much for opening this PR! Looks like a great start already. I pushed the Flax changes in the latest commit. In short, the simplest way of setting the parameters to un-trainable in Flax is by stopping the back-prop through the layers. Otherwise, we need to explicitly pass a dict to the optimiser that defines which parameters are trainable/non-trainable (see https://colab.research.google.com/drive/1K-5bz6R6kt9GAvaUHvzYvvA-IOAO2PhL#scrollTo=BrF6Dtb8GlkJ)

There's not a test to check that the embed params are non-trainable, but you could certainly add one. This could follow the style of test that we use to check that we correctly freeze the encoder when we do decoder-only fine-tuning:

Regarding initialising the weights with sinusoidal embeddings - I agree that this should be the default case! In 99% of cases users will just use the model from pre-trained, in which case the embeddings will be initialised with the sinusoids, but if a user were to randomly initialise the model, the embeddings would be initialised incorrectly.

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

That's a great solution! However, from what I understand, it means that in the Flax implementation, it is not possible (or easy) to re-enable training for positional encodings? (something that we discussed previously)

It's possible (but a bit involved) to add functionality to toggle whether we train the PE's in Flax. However, to me this PR is a bug fix, rather than a feature addition. I agree with what you said in the issue that we should not train the embeddings, since they used fixed sinusoidal embeddings in the original implementation, so I think it's fine if we do a straight fix and always freeze the embeddings here, since this is the correct behaviour.

I added sinusoids weight init for PyTorch implementation. Looking at TF and Flax, I'm not sure where to put weight init. It seems like there is no weight init code in TF? For Flax, I see this but don't really understand what's going on.

From other TF Keras and Flax code I have seen, I think the typical pattern is to pass weight init function to a module when it is created? I'm not sure what is the pattern HF is using here.

The init_weights function in Flax is used to initialise the Flax model's parameters by passing a set of dummy inputs (zeros and ones). Flax traces out the shapes of the weights that you get when you pass these dummy inputs, and initialises weights with the right shapes accordingly (see https://flax.readthedocs.io/en/latest/guides/flax_basics.html#model-parameters-initialization).

This init_weights function won't actually change the values of the weights, just their shapes and dtypes. To change the initialising function, we can pass an argument embedding_init function to the init of the embedding layer: https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.Embed.html#flax.linen.Embed.embedding_init

The init function should be an instance of a JAX initialiser. That is, it should take the PRNG Key as the first argument, as well as the shape and target dtype of the module: https://jax.readthedocs.io/en/latest/jax.nn.initializers.html

@sanchit-gandhi I fixed the embedding init for TF and Flax as you requested. I also add test for TF and Flax. For Flax, I don't add a test for non-trainable sinusoidal embedding, because I don't know how to do it cleanly. For checking the weight init in Flax, I don't know Flax semantics so well, so I added a rather "crude" solution to get the encoder position embeddings.

Very nice @gau-nernst - especially the Flax init which is really clean now 👌 Could the PT init go in _init_weights? Otherwise it all looks good to me!

_init_weights() does not see the module name, it only sees the module itself. To make _init_weights() recognize encoder positional embedding, we probably need to set a private attribute to the module.

Could we change the _init_weights logic to:

1. Loop through all modules
2. Check if module is encoder. If yes: loop through all the sub-modules. When we get to the pos embeddings, do the sinusoidal init
3. Check if module is decoder. If yes: loop through all the sub-modules. When we get to the pos embeddings, do the normal init

I put sinusoidal init in _init_weights() but in a different way. Relying on the fact that nn.Module.apply() will traverse the children in a depth-first search manner (leaf modules will be applied first), if we check for Whisper encoder in _init_weights(), it will override the default initialization for positional embeddings.

    def apply(self: T, fn: Callable[['Module'], None]) -> T:
        ...
        for module in self.children():
            module.apply(fn)
        fn(self)
        return self