flax

cc4279c1 - batch_norm.rst: == should be =