flax

Refactor bridge.Module tests from `wrappers_test.py` to another file.

#4581

Merged