flax

4fa56ad5 - Added mixed precision: computation and weights dtypes