flax

49ef97c3 - Merge pull request #5080 from google:improve-promote-dtype-support