Make FlaxPreTrainedModel a Flax Module
Author: cgarciae
Pull request: no

### Feature request

Issue for discussing #22627 

As stated the idea is to make `FlaxPreTrainedModel` a `nn.Module` so Flax users can easily integrate it into other Flax networks or systems that expect Flax Modules. 

## Requirements
1. Be backward compatible, it is desirable not to brake any existing `transformers` users.
2. When `_do_init=False` the `FlaxPreTrainedModel` should behave like regular Flax Module, this means you cannot use `__call__` directly and must instead use the usual `apply` method.

## Challenges
The main challenge is that Flax Module are dataclasses and use various dataclasses APIs like `dataclasses.replace`. This implies a couple of things:

1. All fields must be specified as class annotations.
2. Most of the current logic in `__init__` has to be done in `__post_init__`.
3. Constructors with `**kwargs` will be tricky to handle.

## Current approach
I've decided solve this problem the following way:
1. `FlaxPreTrainedModel` will define all the needed dataclass fields and to comply with its current signature. `__init__` was refactored into a `__post_init__`.
2. `FlaxPreTrainedModel` sub-classes like `FlaxBertPreTrainedModel` will define its own `__init__` such that they keep their current signature and will then forward everything to `super().__init__`. One thing to note is that sub-classes that inherit from other sub-classes that define `__init__` like `FlaxBertModel(FlaxBertPreTrainedModel)` must define a trivial `__init__` that forwards everything to the parent or else `dataclass` will define one for them according to dataclass semantics.
3. To make `dataclasses.replace` happy, the signature for custom `__init__` functions must accept all the fields names as arguments (i.e. must comply with the dataclass signature) even if it will not use some because e.g. the sub-class will define them on its own before forwarding them to the parent class.
4. I've made the `params` an `Optional` since I'll be expecting it to be `None` when `_do_init=False` (we should not keep the weights aside a Flax Module when its behaving as such).

One thing to note about `3` is that it will require ALL subclasses to define `__init__`, currently some get it for free. We can fix all standard `transformers` models, but user defined sub-classes that reuse `__init__` will break. If this is not good enough we can try to automatically generate a `__inti__` method during `__init_subclass__`.

cc @patrickvonplaten @sanchit-gandhi 

Comments:
One additional note, in the current state of #22479 `dataclasses.replace` still doesn't work, which is why the `test_clone` test is not passing:

https://github.com/huggingface/transformers/pull/22479/files#diff-abf3849ef52688f44671a0752d5a74bf08c861db1eaaab7a4827e52b17cc9dcbR166-R169

> One thing to note about `3` is that it will require ALL subclasses to define `__init__`, currently some get it for free. We can fix all standard `transformers` models, but user defined sub-classes that reuse `__init__` will break. If this is not good enough we can try to automatically generate a `__inti__` method during `__init_subclass__`.

I think this type of users is more advanced and used to seeing breaking changes from time to time. They would typically just pin `transformers` version or update their methods.

`test_clone` has been fixed by fully mimicking the dataclass signature from the custom `__init__` methods.

@sanchit-gandhi do you have some time to look into this? 

Latest PR for the refactor is ongoing: https://github.com/huggingface/transformers/pull/22866

PR remains ongoing - @cgarciae are you able to see this one through to completion? More than happy to discuss any final design decisions and get you another review on the PR as and when required!

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.