Allow loading pretrained shared Pytorch checkpoints into flax models
Motivation: Sharded pytorch checkpoints cannot currently be loaded into flax models; this may be desirable in some cases (e.g. "google/ul2").

Changes: I added an few lines to `modeling_flax_utils.py` to support this behavior. The behavior of the added code exactly matches how sharded checkpoints are loaded in `modeling_utils.py` for pytorch models.

@patrickvonplaten, @patil-suraj
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18170). All of your documentation changes will be reflected on that endpoint.<|||||>Oops! Thanks, just added that import.<|||||>Now you'll need to run `make style` to fix the formatting issues :-)<|||||>done!<|||||>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.