Top-K KL Divergence loss

[AAnoosheh]

What does this PR do?

Type of change: New feature

Overview: Writes a new KLDiv Logits loss which only uses top-k vocab values

Usage

# Add a code snippet demonstrating how to use this

Testing

Before your PR is "Ready for review"

• Make sure you read and follow Contributor guidelines and your commits are signed.
• Is this change backward compatible?: Yes/No
• Did you write any new necessary tests?: Yes/No
• Did you add or update any necessary documentation?: Yes/No
• Did you update Changelog?: Yes/No

Additional Information

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 74.63%. Comparing base (fe52b2a) to head (e7a34bf).
⚠️ Report is 9 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #747      +/-   ##
==========================================
- Coverage   74.68%   74.63%   -0.06%     
==========================================
  Files         192      192              
  Lines       18950    18992      +42     
==========================================
+ Hits        14153    14174      +21     
- Misses       4797     4818      +21     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.