[Flax] Addition of FlaxPegasus

[bhadreshpsavani]

What does this PR do?

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? Please add a link
  to it if that's the case. link of PR
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@patrickvonplaten
@patil-suraj

[bhadreshpsavani]

Hi @patil-suraj and @patrickvonplaten,
I was not able to figure out how to add PegasusSinusoidalPositionalEmbedding in the Flax Version and QuestionAnswering and Classification classes are not added yet since the original torch version don't have the classes. Shall we add it?
Please let me know your review on this PR.

[patil-suraj]

You could find the flax version of SinusoidalPositionalEmbedding in the FlaxMarian

transformers/src/transformers/models/marian/modeling_flax_marian.py

Line 752 in 76c4d8b

self.embed_positions = create_sinusoidal_positions(

Also, Pegasus isn't really intended for QA and classification so it's okay to not add those heads yet.

[patil-suraj]

Thanks a lot for adding this @bhadreshpsavani , great work!

The PR looks good overall, I've left a few comments below.

Mainly

• Let's try to use as many copied from ... statements as possible
• The order of layer_norm in FlaxPegasusEncoder and FlaxPegasusDecoder should be changed. I've left details in the comment.

Let me know if it's not clear or if you need any other help :)

[patrickvonplaten]

Thanks a lot for more or less completing the PR - great job @bhadreshpsavani !

It seems like there are some small differences between the PyTorch & Flax Model. This could be due to slightly different activation functions or small differences with the position ids....

It would be awesome if you could try to debug layer by layer what might be the problem there @bhadreshpsavani

Another possibility is that there is no difference and it's just the framework that causes the difference. In this case, we'll just have to accept it and change the tolerance.

[bhadreshpsavani]

Sure @patrickvonplaten,
I will compare the code and debug it :)

[bhadreshpsavani]

Hi @patil-suraj and @patrickvonplaten,
Thanks for the review and suggestions. Please let me know if anything missing in the PR.

[patil-suraj]

Thanks a lot for fixing the issues, looks good now. If you could give me access to this branch I would like to update the slow tests.

[bhadreshpsavani]

Done!
Please go ahead with the fix for slow tests.
Once this is merged, I will create another PR for that Typo Fix in BART and PEGASUS that I come across

[patil-suraj]

Thanks a lot for adding this model, great work @bhadreshpsavani !

I updated the slow tests and also pushed a couple of flax checkpoints (pegasus-large, pegasus-xum) to the hub. WIll also push the remaining official weights later.

@patrickvonplaten do you wanna give it another look?

[patrickvonplaten]

Awesome!