[Flax] Implement FlaxElectraModel, FlaxElectraForMaskedLM, FlaxElectraForPreTraining
Author: chris-tng
Pull request: yes
Labels: name: WIP

# What does this PR do?

1. Implement Flax version of Electra model : `FlaxElectraModel`, `FlaxElectraForMaskedLM`, `FlaxElectraForPreTraining`. Most of the code taken from FlaxBert version with changes in parameters and forward pass.
2. Adjust `convert_to_pytorch` to load weights for Electra
3. Implement `FlaxElectraGeneratorPredictions`, `FlaxElectraDiscriminatorPredictions` for generator and discriminator prediction head.
4. Implement test in `tests/test_modeling_flax_electra.py`

Forward pass works by running
```shell
pytest tests/test_modeling_flax_electra.py
```

Hi @patrickvonplaten  , @mfuntowicz , I've seen your work on FlaxBert, so I'm tagging in case you want to review. Please note that I use `flax setup` instead of decorator `@nn.compact` since the former
- allows to test and inspect submodule
- separate submodule declaration from the forward pass. Forward pass method can be very long if using `@nn.compact`
I'm happy to revert this change to make code style consistent.

Let me know if you have any questions or feedbacks.
Thanks.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [x] Did you write any new necessary tests? Yes, test added `tests/test_modeling_flax_electra.py`

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors which may be interested in your PR.


Comments:
Hi @chris-tng -- thanks for trying out Flax in HF Transformers!

A quick comment on `nn.compact` and `setup` (I work on Flax) -- indeed if you want access to submodules such as for transfer learning then using `setup` is the way to go. I skimmed the PR and see that you use `setup` in some places and `nn.compact` in others? I'm curious whether you found `nn.compact` more useful in particular settings.

Indeed `setup` is more similar to the PyTorch style (though you still get shape inference if you use `nn.compact` in modules that don't have submodules). `nn.compact` is nice if you want to use loops or conditionals to define submodules based on hyperparameters, and some people also prefer how it "co-locates" their submodule definitions and usage. But ultimate it's somewhat a matter of preference.

(Please do let us know whatever other thoughts or questions on Flax on our discussion board: https://github.com/google/flax/discussions)

Happy holidays and new year!

> shape

Hey @avital,

Thanks a lot for your input here! That's very useful. Most of the main contributors to Transformers are on holiday at the moment and this is a rather big design decision to make going forward with Flax, so I think we'll have to wait here until early January until everybody is back (@sgugger, @LysandreJik, @mfuntowicz) 

Happy holiday to you as well :-) 

Hi @avital ,

Apology for my delayed response. I appreciate your great work on Flax. Regarding the use of `setup()` and `nn.compact`, personally I find `setup` works better for testing submodules. This is useful for converting and debugging the module (and submodules). For instance, I can create a model/module with many submodules:

```python
class Dummy(nn.Module):

  def setup(self):
    self.submodule1 = nn.Dense(10)
    self.submodule2 = MyLayerNorm()

  def __call__(self):
    # do something here
```

After loading model weights from a dict, I can access/debug submodule by simply accessing the attribute: `dummy.submodule1`, `dummy.submodule2`. From this, I can debug forward pass, check model weights of invididual submodule.

Shameless plug, I wrote a blog post about porting huggingface pytorch model to flax, [here](https://chris-tng.github.io/blog/transformer/nlp/flax/pytorch/jax/2020/12/16/tips-flax.html). I'm a new Flax user so please correct me if I'm missing anything.

Happy holiday and happy new year to everyone 🎄 🍾 

Hey @chris-tng, 

sorry to had you wait for this long. I'll solve the merge conflicts in your PR and then use your PR to change the `@nn.compact` to `setup` in all other flax models as well so that we have a common standard now. Since most of our users are used to the "PyTorch" style and I only see advantages for our library philosophy:

- We base most design decisions on the PyTorch style
- We prefer slightly less compact readable code over the slightly more "magic" functionalities that might reduce code
- To me the are no real downsides to using `setup`

Intermediate state is saved here: #9484 will push to this PR on Monday the latest

Hey @chris-tng,

I noticed that we will probably have to wait a bit to get this merged: https://github.com/google/flax/pull/683 to be able to continue the PR. Will keep you up-to-date :-)

Hi folks, sorry for the delay with the new-year shuffle and school shutdown.

google/flax#683 required a bit more conversation and updating some other codebases but now it's merged! If you have a moment, please take a look and see if it helps unblock progress. We'll release Flax 0.4.0 soon, but installing from GitHub now is the way to go.

Hey, sorry for barging in
I was needing a small BERT-like model in Jax and so I've recently updated this in a local branch to work with the Flax refactoring that makes checkpoints directly compatible with PyTorch (plus fixing some other issues that had gone through the cracks)
Should I push directly to this branch or make a new PR from my fork? Also should I wait #11364 and update my code accordingly?

> Hey, sorry for barging in
> I was needing a small BERT-like model in Jax and so I've recently updated this in a local branch to work with the Flax refactoring that makes checkpoints directly compatible with PyTorch (plus fixing some other issues that had gone through the cracks)
> Should I push directly to this branch or make a new PR from my fork? Also should I wait #11364 and update my code accordingly?

Hey @CoderPat,

It would be great if you could wait until #11364 is merged (should be done in the next 2 days). The PR fixes a couple of bugs :-)

No problem @patrickvonplaten! Also regarding git logistics, is it better to ask @chris-tng for permission to push directly to his branch?

> No problem @patrickvonplaten! Also regarding git logistics, is it better to ask @chris-tng for permission to push directly to his branch?

I think it's alright to copy past the code that is still useful and open a new branch, if you'd like to add Electra :-). On the branch we should then give credit to @chris-tng , but since the PR is quite old now I think he would be fine if we close this one and open a new one (Please let me know if this is not the case @chris-tng :-)) . #11364 should be the last refactor before the "fundamental" Flax design is finished.

Just to confirm @patrickvonplaten , the flax refactor is merged and the structure should be stable enough that I can work on implementing Electra right?

Exactly @CoderPat - very much looking forward to your PR :-)

Closing as this PR is super old and partly fixed by #11426