flax

cdf3b0c3 - Added mixed precision: computation and weights dtypes