transformers

[Flax] adding support for batch norm layers

#21581

Merged