flax
9e6b8193 - Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.

Commit

1 year ago

Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful. Fixes https://github.com/google/flax/issues/4008 PiperOrigin-RevId: 646679331

References

#4029 - Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.

Author

IvyZX

Committer

a-googler

Parents

3b218708

flax 9e6b8193 - Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.

flax
9e6b8193 - Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.