Resnet flax
# Flax Implementation of `microsoft/resnet-50`

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

@sanchit-gandhi  I guess :sweat_smile: 

## Status

Last Updated - Sunday, 12 February 2023

### TODO
 - [x] Blocked on merge of [this PR](https://github.com/huggingface/transformers/pull/21581) to add support for BatchNorm layers. 
 - [ ] Uploading [Shubhamai/resnet-50](https://huggingface.co/Shubhamai/resnet-50) flax weights to [microsoft/resnet-50](https://huggingface.co/microsoft/resnet-50).
_The documentation is not available anymore as the PR was closed or merged._<|||||>Very cool! Sorry to only reply here now - looks like you've made a really solid start to this PR! Let's get the Batch Norm PR merged ASAP and then go full send on Flax ResNet! 🚀

Feel free to ping me with any questions / queries! More than happy to help with the integration!<|||||>@sanchit-gandhi the PR is also now ready for your review, thanks a lot for your time. <|||||>Thanks for the review, I will make the changes soon. Many of the reviews here also apply to the [Flax Convnext PR](https://github.com/huggingface/transformers/pull/21485) so will make corresponding changes there too. <|||||>Hey @Shubhamai! Really nice work on this PR - just a few small changes to go now! Feel free to ping me once you're happy with the last bits and I'll get you a final review. Taking a look at Flax RegNet in the meantime!<|||||>Made the request changes and apologies for the late reply, had a busy schedule lately.<|||||>Thanks for the review @amyeroberts 🙌 Feel free to propagate the changes forward to [Flax RegNet](https://github.com/huggingface/transformers/pull/21867) @Shubhamai  and we can get a final review there!<|||||>We can merge when the CI is green! 🟢<|||||>Thank you so much for the review & merge @sanchit-gandhi @amyeroberts , really appreciate your work & time on this ❤️ 