[Flax] Big FlaxBert Refactor (#11364)
* improve flax
* refactor
* typos
* Update src/transformers/modeling_flax_utils.py
* Apply suggestions from code review
* Update src/transformers/modeling_flax_utils.py
* fix typo
* improve error tolerance
* typo
* correct nasty saving bug
* fix from pretrained
* correct tree map
* add note
* correct weight tying
