[Flax] improve large model init and loading 
Author: patil-suraj
Pull request: yes

# What does this PR do?

As discussed in [#15766](https://github.com/huggingface/transformers/issues/15766#issuecomment-1047912290) this PR adds the `_do_init` argument to handle large model loading in flax. By default `_do_init=True` and the API stays same.

1. `__init__`:

- When `_do_init=False` is passed to `__init__`, the params are not initialised and the user should manually call the `model.init_weights` method to do the initialisation. 
- When `_do_init` is `False` accessing `model.params` is not allowed and params must be always kept outside of the model
- This PR also adds the `params_shape_tree`  shape property to `FlaxPreTrainedModel`, which is a PyTree with shape and dtype information for each param.

This is how the API looks like:
```python
config = BertConfig()
model = FlaxBertModel(config, _do_init=False)

# accessing model.params will raise an ValueError
model.params 

# to init the params
params = model.init_weights(model.key, model.input_shape)

# setting the model.params will also raise an ValueError
model.params = params

# model.init_weights can be used with `jit` and `pjit` to init the params in CPU or init in sharded way
params = jax.jit(model.init_weights, static_argnums=1, backend="cpu")(model.key, model.input_shape)
```

2. `from_pretrained`:

- When `_do_init=False` is passed to `from_pretrained`, the model won't be randomly initialised, only the weights will be loaded and returned  along with the model instance.
- when `_do_init=False` params will always be loaded on CPU
- If `_do_init=False` and some keys are missing, the user should call `init_weights` and pass it the loaded params. `init_weights` will then take care of adding the missing keys.
- as described above, getting and setting `model.params` will raise an error.

```python3
model, params = FlaxBertModel.from_pretrained("...", _do_init=False)
# if keys are missing
params = model.initt_weights(model.key, model.input_shape, params)
```

cc @borisdayma @sanchit-gandhi  for info.

Fixes #15766 

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Hi @patil-suraj,

Thanks a lot for fixing this issue.
Could you please give us some estimation of when this pull request will be ready?

Hey @agemagician !
Thanks. I just need to run and fix a few tests and then it should be good to merge by tomorrow.

> Hey @agemagician ! Thanks. I just need to run and fix a few tests and then it should be good to merge by tomorrow.

Perfect, thanks a lot @patil-suraj for your effort.
It will be awesome if you could update one of the flax language model examples to understand the changes that is needed after this merge:
https://github.com/huggingface/transformers/tree/main/examples/flax/language-modeling

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.

Hi @patil-suraj ,

Any updates regarding this merge ?

@patil-suraj - think we can merge this no?

Just need to fix the template tests and then we can merge it.