Fix sample_logits crash when top_k exceeds vocab size

[robbiebusinessacc]

sample_logits passes top_k straight to final_logits.topk(top_k). When
top_k is larger than the vocabulary size (e.g.
model.generate(top_k=100_000) on a small-vocab model), .topk() raises
RuntimeError: selected index k out of range.

HuggingFace's TopKLogitsWarper handles this by clamping (top_k = min(top_k, logits.size(-1))). This change does the same: after the existing top_k > 0
assertion, clamp top_k to the last-dim size before calling .topk().
Behaviour is unchanged whenever top_k <= vocab_size.

While here, I removed a stale #! TODO: Finish testing all the edge cases
scratch-code comment and replaced it with a one-line docstring note describing
the clamping behaviour, since this PR adds those edge-case tests.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature
  works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect
  backward compatibility

Notes

sample_logits had no test coverage before this. The PR adds tests for the
top_k > vocab_size clamp (the crash case), exact-vocab-size top_k, plus
existing-behaviour regression cases. Verified by reverting only the clamp
line, which reproduces the original RuntimeError: selected index k out of range; restoring it makes all tests pass. black, isort, and pycln are
clean.

[jlarson4]

Excellent work on this! Merging for inclusion in the next release