TTSKit : guard sampleFromProbs against zero-sum top-k probabilities (#450)

[achyutbenz19]

Summary

Fixes #450.

GreedyTokenSampler.sampleFromProbs crashes with Fatal error: Can't get random value with an empty range when every top-k probability rounds to zero. The call is Float.random(in: 0..<probSum, using: &rng); when probSum == 0 the range is empty and Swift traps.

The reporter triggers this during long-form TTS generation (audiobook chapters) at low temperature (0.10) and small top-k (15). Numerical underflow accumulates across hundreds of chunks until a chunk's softmax output lands entirely below single-precision range for those 15 slots.

Scope of the change

Sources/TTSKit/Utilities/Sampling.swift, +12/-0 on one method.

1. In the top-k branch, guard probSum > 0 before calling Float.random(in: 0..<probSum). On underflow, fall back to greedy selection by returning the highest-probability token. MLTensor.topK returns entries in descending probability order, so idxArray.first is the argmax.
2. In the full-distribution branch (no top-k), apply the same guard on probsArray.reduce(0, +) for completeness. Softmax normally sums to 1, but the same numerical path can underflow in principle.

The guard changes behavior only when probSum == 0, which is the crash case. All paths where the old code returned a valid token continue to do the same random draw bit-for-bit.

Match to reporter's proposal

The reporter suggested:

Guard against a zero sum and fall back to the highest-probability token (greedy)

That is exactly what this patch does. Greedy fallback (rather than, say, uniform-random over top-k) matches the intent of a temperature that has already collapsed the distribution.

Build

Builds clean on Xcode 26.1.1 / Swift 6.2.1 against the TTSKit product target.

What this does not do

• Does not change non-underflow behavior. Every codepath where probSum > 0 is byte-identical, including RNG advancement order.
• Does not change the caller-side contract. sampleFromProbs still returns Int32.
• Does not touch the MLTensor-from-logits sampler a few lines below, since its underflow profile is different (it draws from a pre-flattened array) and the reporter did not flag it.

Tools used

git, swift build, and audiokit for the rest of the audio PRs in this session. Skipped running an audiokit differential matrix on this specific change because the patch only adds guards along the crash-only path (probSum == 0); every branch where probSum > 0 is byte-for-byte identical to the pre-patch code, including RNG draw order. No audio-output differential is possible against a path that crashes pre-patch.

Disclosure

I am an AI assistant (Anthropic's Claude) helping a user contribute this fix. I reproduced the crash path by inspection of the code and verified the patch compiles against current main.