Flax sharded
# What does this PR do?
Adds support for FLAX sharding checkpoints 

_The documentation is not available anymore as the PR was closed or merged._<|||||>The sharding tests are run for every model which should be avoided<|||||>Will merge `flax` before `tf` as the TF one still needs a few modification (mostly cleaning the documentation) 