flax

5e82ac8d - Merge pull request #4920 from google:add-preferred_element_type-arg