Match Flax batch norm params to PyTorch defaults

[martinkim0]

See https://flax.readthedocs.io/en/latest/api_reference/flax.linen/layers.html#flax.linen.BatchNorm and https://pytorch.org/docs/stable/generated/torch.nn.BatchNorm1d.html#torch.nn.BatchNorm1d

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (b6b7b86) 89.27% compared to head (b85ba53) 89.27%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2309   +/-   ##
=======================================
  Coverage   89.27%   89.27%           
=======================================
  Files         145      145           
  Lines       11827    11827           
=======================================
  Hits        10559    10559           
  Misses       1268     1268           

Files	Coverage Δ
scvi/module/_jaxvae.py	100.00% <100.00%> (ø)

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.