Skip to content

TTSKit: sampleFromProbs crashes when top-k probabilities sum to zero #450

Description

@ddegner

Description

GreedyTokenSampler.sampleFromProbs in Sources/TTSKit/Utilities/Sampling.swift crashes with:

Fatal error: Can't get random value with an empty range

when all top-k probabilities collapse to zero after softmax, making probSum == 0 and Float.random(in: 0..<0) an empty range.

Steps to reproduce

Run TTS generation on long-form text (e.g. audiobook chapters) with temp=0.10, topK=15. After hundreds of chunks, softmax occasionally produces a distribution where all top-k values round to Float zero.

Location

https://github.com/argmaxinc/WhisperKit/blob/d447d30/Sources/TTSKit/Utilities/Sampling.swift#L238-L243

let probSum = probsArray.reduce(0, +)
let randomValue = Float.random(in: 0..<probSum, using: &rng)  // 💥 when probSum == 0

Suggested fix

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

let probSum = probsArray.reduce(0, +)
guard probSum > 0 else {
    return Int32(idxArray[0])  // topK is already sorted descending
}
let randomValue = Float.random(in: 0..<probSum, using: &rng)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions