Add _apply_decay_mask_logits function to avoid numerical instability

[boxin-wbx]

No description provided.

[timoschick]

Thanks for this pull request 👍
Is there a particular reason for not actually replacing calls to _apply_decay_mask with calls to the new _apply_decay_mask_logits (and deleting the former function altogether)?

[boxin-wbx]

Thanks for the comment.

I think we can keep both versions. If there is no numerical instability issue, we can of course use _apply_decay_mask. Also, if you have time, you can run both versions and have a comparison between these two implementations. I suppose these two versions should not differ much.