jax

192e9086 - Merge pull request #2561 from madisonmay/batch-norm-no-scale-or-center