huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-15 02:16:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix the gumbel softmax by casting to f32. (#2928)

Browse Source
This commit is contained in:
Laurent Mazare
2025-04-28 19:48:51 +02:00
committed by GitHub
parent e98754fc5a
commit d4bac37a61
1 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
candle-nn/src/sampling.rs
View File

@ -8,13 +8,16 @@ pub fn gumbel_softmax<D: candle::shape::Dim>(
) -> Result<Tensor> {
if temperature <= 0.0 {
logits.argmax(dim)
} else if temperature == 1.0 {
let minus_g = logits.rand_like(1e-7, 0.999)?.log()?.neg()?.log()?;
let sampled = (logits - minus_g)?.argmax(dim)?;
Ok(sampled)
} else {
// Cast to f32, doing the Gumbel softmax in bf16 is a bit unstable.
let logits = logits.to_dtype(candle::DType::F32)?;
let minus_g = logits.rand_like(1e-7, 0.999)?.log()?.neg()?.log()?;
let sampled = (logits + minus_g * (-temperature))?.argmax(dim)?;
Ok(sampled)
if temperature == 1.0 {
let sampled = (logits - minus_g)?.argmax(dim)?;
Ok(sampled)
} else {
let sampled = (logits + minus_g * (-temperature))?.argmax(dim)?;
Ok(sampled)
}
}
}