flax
Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.
#4029
Merged
Merged