Fixes `random_flax_module` with `flax.linen.BatchNorm`#1823

Merged

fehiepsi merged 8 commits intopyro-ppl:masterfrom juanitorduz:nn_batch_issue_1446

Jul 1, 2024

+24-3

Commits on Jun 19, 2024

filter oout tests waiting for next tfp release
juanitorduz
committed

Commits on Jun 30, 2024