flax

27fd2284 - cap dynamic scale to float32 max