flax

7c36812f - Merge pull request #4613 from google:fix-promote-dtypes