Flax dtype-dependent numerical masking
# What does this PR do?

Fixes #21176

For some models, our Flax numerical masking was incompatible with the desired variable type. This PR fixes it by selecting a numerical mask that is the minimum for the corresponding variable type.

This PR is akin to #17306 for PT. Thank you @LysandreJik and @ydshieh for pointing it out 🙏 
@sgugger this solution was discussed on Slack with the Flax team, hence no added Flax reviewers :)<|||||>_The documentation is not available anymore as the PR was closed or merged._