PR #21581 [Flax] adding support for batch norm layers

[Flax] adding support for batch norm layers #21581

sgugger merged 7 commits into huggingface:main from Shubhamai:flax-batchnorm-support

[flax] adding support for batch norm layers

e4ae180f

fixing bugs related to pt+flax integration

405b8479

Merge branch 'huggingface:main' into flax-batchnorm-support

b0ead9ac

cleanup, batchnorm support in sharded pt to flax

a6a5fc06

support for batchnorm tests in pt+flax integration

4bf496ca

Shubhamai changed the title ~~[WIP] [Flax] adding support for batch norm layers~~ [Flax] adding support for batch norm layers 2 years ago

Shubhamai marked this pull request as ready for review 2 years ago

sanchit-gandhi commented on 2023-02-22

Merge branch 'main' into flax-batchnorm-support

6397800c

simplifying checking batch norm layer

1c2e91a8

Shubhamai requested a review from

sanchit-gandhi 2 years ago

sanchit-gandhi approved these changes on 2023-02-23

sanchit-gandhi requested a review from

sgugger 2 years ago

sgugger approved these changes on 2023-02-24

sgugger merged f7ca656f into main 2 years ago

Reviewers

sgugger

sanchit-gandhi

Assignees

No one assigned

Labels

None yet

Milestone

No milestone