[Flax] adding support for batch norm layers#21581

sgugger merged 7 commits into

huggingface:mainhuggingface/transformers:mainfrom

Shubhamai:flax-batchnorm-supportShubhamai/transformers:flax-batchnorm-supportCopy head branch name to clipboard