core/go
8
0
Fork 0
generated from lthn/LEM
Code Issues 2 Pull requests 1 Projects Releases Packages Wiki Activity Actions

fix: remove unused vars in TopP sampler placeholder

Browse source 
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Claude 2026-02-16 01:54:29 +00:00 • committed by Snider
parent a0f77960a1
commit e9d9a3c3a0
1 changed files with 3 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
pkg/mlx/sample/sample.go
View file

@ -75,24 +75,9 @@ func (k TopKSampler) Sample(logits *mlx.Array) *mlx.Array {
type TopP float32 type TopP float32
func (p TopP) Sample(logits *mlx.Array) *mlx.Array { func (p TopP) Sample(logits *mlx.Array) *mlx.Array {
// Softmax to get probabilities // TODO: full nucleus sampling requires cumsum which mlx-c doesn't expose directly.
probs := mlx.Softmax(logits) // For now, pass through. TopK + Temperature covers most use cases.
// Sort descending return logits
neg := mlx.Negative(probs)
sortedIdx := mlx.Argpartition(neg, 0, -1)
sortedProbs := mlx.Take(probs, sortedIdx, -1)
// Cumulative sum
cumProbs := mlx.Sum(sortedProbs, -1, true) // simplified — full impl needs cumsum
// Mask tokens beyond threshold
threshold := mlx.FromValue(float32(p))
mask := mlx.Where(
mlx.FromValue(true), // placeholder — proper impl compares cumprobs > p
mlx.FromValue(float32(math.Inf(-1))),
logits,
)
return mask
} }
// MinPSampler masks tokens below min_p * max_prob. // MinPSampler masks tokens below min_p * max_prob.