[Flax] adding support for batch norm layers
Author: Shubhamai
Pull request: yes

# What does this PR do?

Adding support for batch norm layers in flax. Several models like ResNets, MobileNets have BatchNorm layers which (for now) cannot be ported to HuggingFace Flax due to Flax requiring an extra `batch_stats` key along with `params` to perform training and inference. 

The changes here are tested locally with my work-in-progress resnet PR [https://github.com/huggingface/transformers/pull/21472](https://github.com/huggingface/transformers/pull/21472), ~~though more testing is needed.~~ LGTM :rocket: 

### Design Choices
_Mentioning a few design decisions and some notes to ease the review process_ :)

- To detect whether the flax model contains batch norm layers, we use a simple if statement 
        `if "params" in flax_model.params and "batch_stats" in flax_model.params`
        An alternative could have been adding a `contains_batchnorm` attribute to `FlaxPreTrainedModel` class but I decided not to go with it since found the first option simpler and kept no. of class attributes as low as possible unless strictly needed. 
- `num_batches_tracked` is not transferred from pt to flax weight conversion and are newly initialized when doing flax to pt weight conversion.
- ~~`pt_flax_cross_test` in `test_modeling_flax_common.py` currently does not support comparing the weights between  `batch_stats` and PyTorch equivalent BatchNorm mean and var. I am still currently looking into supporting it.~~ 
        It does now :)

## 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.
- [ ] 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).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

- Flax: @sanchit-gandhi

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Hi @sanchit-gandhi, the PR is now ready for your review. I'm sorry if my previous ping was confusing since I mentioned it when the PR was in WIP. Next time, I will make sure to ping you only when it's ready for review. Thank you so much for taking the time to review the PR.

[Flax Convnext PR](https://github.com/huggingface/transformers/pull/21485) is also ready for your review. 

HI, @sgugger  I have been trying to mention @sanchit-gandhi but unfortunately have not received a response. I was hoping you could assist me in finding the right person to reach out to for a review. Any help would be greatly appreciated.

The right person is @sanchit-gandhi 

Cool! Such a pleasure to work with :hugs:  team & transformers. Along with ResNets, I have tested this PR on flax implementation of RegNet and EfficientNet ( which uses batch norm ) locally. Both of which are nearly complete and will make PR once this get's merged.

Super fun working with you too @Shubhamai! Thank you for your contribution and high quality PR! Very excited to see the follow-ups with ResNet, RegNet and EfficientNet, these will be very nice additions to the library. Feel free to ping me when they're in a review ready state and I'll take a look as quickly as possible! Likewise if you have any Flax questions / queries I'm more than happy to help!