Shard checkpoint for `tf` and `flax`
### Feature request

The same sharding capabilities as pytorch should be available to `flax` and `tf`. This is required in order to push the OPT30B model. 

### Motivation

Pushing $>45GB$ models (and have the same behaviour as in pytorch ). 

### Your contribution

Could start working on that when I will be back from holidays! 