ronantakizawa/codereview-bench huggingface dataset:
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
package sample

import (
	"math"
	"slices"
)

// apply temperature scaling and softmax to the logits
func temperature(ts []token, temp float32) []token {
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Apply temperature and compute exp(x - max)
	temp = max(temp, 1e-7)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
		sum += ts[i].value
	}

	// Normalize
	for i := range ts {
		ts[i].value /= sum
	}

	return ts
}


Review comment (suggestion):
GoDoc this

```
// temperature applies sacling and softmax to the logits
```

Code after:
package sample

import (
	"math"
	"slices"
)

// temperature applies scaling and softmax to the logits
func temperature(ts []token, temp float32) []token {
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Apply temperature and compute exp(x - max)
	temp = max(temp, 1e-7)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
		sum += ts[i].value
	}

	// Normalize
	for i := range ts {
		ts[i].value /= sum
	}

	return ts
}
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
		return ts
	}

	maxProb := float32(math.Inf(-1))
	for _, token := range ts {
		if token.value > maxProb {
			maxProb = token.value
		}
	}

	threshold := maxProb * float32(p)

	// Filter tokens in-place
	validTokens := ts[:0]
	for i, token := range ts {
		if token.value >= threshold {
			validTokens = append(validTokens, ts[i])
		}
	}

	ts = validTokens
	return ts
}

// TODO(parthsareen): possibly replace with simpler implementation https://github.com/ollama/ollama/issues/9584
// counting sort implementation to sort tokens by logits
func sortLogits(tokens []token) {
	if len(tokens) <= 1 {
		return
	}

	// Find max/min in a single pass
	minLogit, maxLogit := tokens[0].value, tokens[0].value
	for _, t := range tokens[1:] {
		if t.value < minLogit {
			minLogit = t.value
		} else if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Calculate scaling to map to uint32 range
	logitRange := maxLogit - minLogit
	if logitRange < 1e-6 {
		return // All values effectively equal
	}

	// Count frequencies directly from tokens
	const maxInt = (1 << 24) - 1 // Use 24 bits for good granularity
	var counts [256]int          // For first byte


Review comment (suggestion):
```suggestion
// sortLogits sorts implementation to sort tokens by logits using counting sort...
```

I'd include why counting sort over built in

Code after:
		return ts
	}

	maxProb := float32(math.Inf(-1))
	for _, token := range ts {
		if token.value > maxProb {
			maxProb = token.value
		}
	}

	threshold := maxProb * float32(p)

	// Filter tokens in-place
	validTokens := ts[:0]
	for i, token := range ts {
		if token.value >= threshold {
			validTokens = append(validTokens, ts[i])
		}
	}

	ts = validTokens
	return ts
}

// TODO(parthsareen): possibly replace with simpler implementation https://github.com/ollama/ollama/issues/9584
// sortLogits sorts implementation to sort tokens by logits using counting sort
// counting sort is faster than built-in sort for this use case
func sortLogits(tokens []token) {
	if len(tokens) <= 1 {
		return
	}

	// Find max/min in a single pass
	minLogit, maxLogit := tokens[0].value, tokens[0].value
	for _, t := range tokens[1:] {
		if t.value < minLogit {
			minLogit = t.value
		} else if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Calculate scaling to map to uint32 range
	logitRange := maxLogit - minLogit
	if logitRange < 1e-6 {
		return // All values effectively equal
	}

	// Count frequencies directly from tokens
	const maxInt = (1 << 24) - 1 // Use 24 bits for good granularity
	var counts [256]int          // For first byte
Repository: ollama/ollama
File: server/routes.go (Go)

Code before:
	delete(kvData, "tokenizer.chat_template")
	resp.ModelInfo = kvData

	tensorData := make([]api.Tensor, len(tensors.Items()))
	for cnt, t := range tensors.Items() {
		tensorData[cnt] = api.Tensor{Name: t.Name, Type: t.Type(), Shape: t.Shape}
	}
	resp.Tensors = tensorData

	if len(m.ProjectorPaths) > 0 {
		projectorData, err := getKVData(m.ProjectorPaths[0], req.Verbose)
		if err != nil {
			return nil, err
		}
		resp.ProjectorInfo = projectorData
	}

	return resp, nil
}

func getModelData(digest string, verbose bool) (ggml.KV, ggml.Tensors, error) {
	maxArraySize := 0
	if verbose {
		maxArraySize = -1
	}
	data, err := llm.LoadModel(digest, maxArraySize)
	if err != nil {
		return nil, ggml.Tensors{}, err
	}

	kv := data.KV()

	if !verbose {
		for k := range kv {
			if t, ok := kv[k].([]any); len(t) > 5 && ok {
				kv[k] = []any{}
			}
		}
	}

	return kv, data.Tensors(), nil
}

func getKVData(digest string, verbose bool) (ggml.KV, error) {
	kv, _, err := getModelData(digest, verbose)
	return kv, err
}

func (s *Server) ListHandler(c *gin.Context) {
	ms, err := Manifests(true)
	if err != nil {

Review comment (suggestion):
This was already here, so it can be a follow up, but does this load the whole model into memory? The function name and comment implies this, feels like we might be misusing this function if so.

Code after:
	delete(kvData, "tokenizer.chat_template")
	resp.ModelInfo = kvData

	tensorData := make([]api.Tensor, len(tensors.Items()))
	for cnt, t := range tensors.Items() {
		tensorData[cnt] = api.Tensor{Name: t.Name, Type: t.Type(), Shape: t.Shape}
	}
	resp.Tensors = tensorData

	if len(m.ProjectorPaths) > 0 {
		projectorData, _, err := getModelData(m.ProjectorPaths[0], req.Verbose)
		if err != nil {
			return nil, err
		}
		resp.ProjectorInfo = projectorData
	}

	return resp, nil
}

func getModelData(digest string, verbose bool) (ggml.KV, ggml.Tensors, error) {
	maxArraySize := 0
	if verbose {
		maxArraySize = -1
	}
	data, err := llm.LoadModel(digest, maxArraySize)
	if err != nil {
		return nil, ggml.Tensors{}, err
	}

	kv := data.KV()

	if !verbose {
		for k := range kv {
			if t, ok := kv[k].([]any); len(t) > 5 && ok {
				kv[k] = []any{}
			}
		}
	}

	return kv, data.Tensors(), nil
}

func (s *Server) ListHandler(c *gin.Context) {
	ms, err := Manifests(true)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	models := []api.ListModelResponse{}
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
package sample

import (
	"container/heap"
	"math"
	"slices"
)

// tokenHeap implements heap.Interface and holds tokens as a min-heap to track k largest elements
type tokenHeap []token

func (h tokenHeap) Len() int           { return len(h) }
func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value } // Use < for min-heap to track largest elements
func (h tokenHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *tokenHeap) Push(x any) {
	*h = append(*h, x.(token))
}

func (h *tokenHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// temperature applies scaling and softmax to the logits
func temperature(ts []token, temp float32) []token {
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Apply temperature and compute exp(x - max)

Review comment (suggestion):
```suggestion
func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value }
```

self explanatory

Code after:
package sample

import (
	"container/heap"
	"math"
	"slices"
)

// tokenHeap implements heap.Interface and holds tokens as a min-heap to track k largest elements
type tokenHeap []token

func (h tokenHeap) Len() int           { return len(h) }
func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h tokenHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *tokenHeap) Push(x any) {
	*h = append(*h, x.(token))
}

func (h *tokenHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// temperature applies scaling and softmax to the logits
func temperature(ts []token, temp float32) []token {
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Apply temperature and compute exp(x - max)
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Apply temperature and compute exp(x - max)
	temp = max(temp, 1e-7)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
		sum += ts[i].value
	}

	// Normalize
	for i := range ts {
		ts[i].value /= sum
	}

	return ts
}

// topK limits the number of tokens considered to the k highest logits
func topK(ts []token, k int) []token {
	if k >= len(ts) {
		sortLogits(ts)
		return ts
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, ts[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(ts); i++ {
		if ts[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, ts[i])
		}
	}

	// Convert heap to sorted slice in descending order
	result := make([]token, k)
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(&h).(token)
	}

	return result

Review comment (suggestion):
This is a TODO for later, but it's not clear from the function signature here whether `ts` is mutated or not. The return value of `[]token` implies it isn't, yet I think it is. We'll need to clear that up.

Code after:
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Apply temperature and compute exp(x - max)
	temp = max(temp, 1e-7)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
		sum += ts[i].value
	}

	// Normalize
	for i := range ts {
		ts[i].value /= sum
	}

	return ts
}

// topK limits the number of tokens considered to the k highest logits
func topK(ts []token, k int) []token {
	if k >= len(ts) || k <= 0 {
		slices.SortFunc(ts, func(a, b token) int {
			switch {
			case a.value < b.value:
				return 1
			case a.value > b.value:
				return -1
			default:
				return 0
			}
		})
		return ts
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, ts[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(ts); i++ {
		if ts[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, ts[i])
		}
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Apply temperature and compute exp(x - max)
	temp = max(temp, 1e-7)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
		sum += ts[i].value
	}

	// Normalize
	for i := range ts {
		ts[i].value /= sum
	}

	return ts
}

// topK limits the number of tokens considered to the k highest logits
func topK(ts []token, k int) []token {
	if k >= len(ts) {
		sortLogits(ts)
		return ts
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, ts[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(ts); i++ {
		if ts[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, ts[i])
		}
	}

	// Convert heap to sorted slice in descending order
	result := make([]token, k)
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(&h).(token)
	}

	return result
}


Review comment (refactor):
Given how rare this is, why don't we just use built-in sort for now. It will simplify the code a lot:

```
	if k >= len(ts) {
		slices.SortFunc(ts, func(a, b token) int {
			switch {
			case a < b:
				return -1
			case a > b:
				return 1
			default:
				return 0
			}
		})
		return ts
	}
```

Code after:
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Apply temperature and compute exp(x - max)
	temp = max(temp, 1e-7)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64((v.value - maxLogit) / temp)))
		sum += ts[i].value
	}

	// Normalize
	for i := range ts {
		ts[i].value /= sum
	}

	return ts
}

// topK limits the number of tokens considered to the k highest logits
func topK(ts []token, k int) []token {
	if k >= len(ts) || k <= 0 {
		slices.SortFunc(ts, func(a, b token) int {
			switch {
			case a.value < b.value:
				return 1
			case a.value > b.value:
				return -1
			default:
				return 0
			}
		})
		return ts
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, ts[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(ts); i++ {
		if ts[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, ts[i])
		}
	}
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:

	return ts
}

// topK limits the number of tokens considered to the k highest logits
func topK(ts []token, k int) []token {
	if k >= len(ts) {
		sortLogits(ts)
		return ts
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, ts[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(ts); i++ {
		if ts[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, ts[i])
		}
	}

	// Convert heap to sorted slice in descending order
	result := make([]token, k)
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(&h).(token)
	}

	return result
}

// topP limits tokens to those with cumulative probability p
func topP(ts []token, p float32) []token {
	if p == 1.0 {
		return ts
	}

	// Find cutoff index where cumulative sum exceeds p
	var sum float32
	for i, t := range ts {
		sum += t.value
		if sum > float32(p) {
			ts = ts[:i+1]
			return ts
		}
	}

	return ts
}

Review comment (suggestion):
```suggestion
	result := make([]token, len(h))
```

Code after:

	return ts
}

// topK limits the number of tokens considered to the k highest logits
func topK(ts []token, k int) []token {
	if k >= len(ts) || k <= 0 {
		slices.SortFunc(ts, func(a, b token) int {
			switch {
			case a.value < b.value:
				return 1
			case a.value > b.value:
				return -1
			default:
				return 0
			}
		})
		return ts
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, ts[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(ts); i++ {
		if ts[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, ts[i])
		}
	}

	// Convert heap to sorted slice in descending order
	result := make([]token, len(h))
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(&h).(token)
	}

	return result
}

// topP limits tokens to those with cumulative probability p
func topP(ts []token, p float32) []token {
	if p == 1.0 {
		return ts
	}

	// Find cutoff index where cumulative sum exceeds p
	var sum float32
	for i, t := range ts {
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
	if k >= len(ts) {
		sortLogits(ts)
		return ts
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, ts[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(ts); i++ {
		if ts[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, ts[i])
		}
	}

	// Convert heap to sorted slice in descending order
	result := make([]token, k)
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(&h).(token)
	}

	return result
}

// topP limits tokens to those with cumulative probability p
func topP(ts []token, p float32) []token {
	if p == 1.0 {
		return ts
	}

	// Find cutoff index where cumulative sum exceeds p
	var sum float32
	for i, t := range ts {
		sum += t.value
		if sum > float32(p) {
			ts = ts[:i+1]
			return ts
		}
	}

	return ts
}

// minP limits tokens to those with cumulative probability p
func minP(ts []token, p float32) []token {
	if p == 1.0 {
		return ts
	}

Review comment (performance):
This can be simplified to:

```

	h := make(tokenHeap, 0, k)

	for _, t := range ts {
		if len(h) < k {
			heap.Push(&h, t)
			continue
		}
		if t.value > h[0].value {
			h[0] = t
			heap.Fix(&h, 0)
		}
	}
```

Double check that the performance doesn't worsen (I think it should improve)

Code after:
	if k >= len(ts) || k <= 0 {
		slices.SortFunc(ts, func(a, b token) int {
			switch {
			case a.value < b.value:
				return 1
			case a.value > b.value:
				return -1
			default:
				return 0
			}
		})
		return ts
	}

	// Initialize min-heap with first k elements
	h := make(tokenHeap, k)
	copy(h, ts[:k])
	heap.Init(&h)

	// Process remaining elements
	for i := k; i < len(ts); i++ {
		if ts[i].value > h[0].value {
			heap.Pop(&h)
			heap.Push(&h, ts[i])
		}
	}

	// Convert heap to sorted slice in descending order
	result := make([]token, len(h))
	for i := k - 1; i >= 0; i-- {
		result[i] = heap.Pop(&h).(token)
	}

	return result
}

// topP limits tokens to those with cumulative probability p
func topP(ts []token, p float32) []token {
	if p == 1.0 {
		return ts
	}

	// Find cutoff index where cumulative sum exceeds p
	var sum float32
	for i, t := range ts {
		sum += t.value
		if sum > float32(p) {
			ts = ts[:i+1]
			return ts
		}
	}
Repository: ollama/ollama
File: fs/ggml/ggml.go (Go)

Code before:
				partialOffload,
				4*batch*(1+
					2*embedding+
					embeddingHeadsK*heads+
					context+
					context*heads)+
					4*embeddingHeadsK*context+
					4*context*embeddingHeadsK+
					4*qkvBias.Shape[0],
			)
		}
	}

	return
}

func (llm GGML) VisionGraphSize() (weights, graphSize uint64) {
	for name, layer := range llm.Tensors().GroupLayers() {
		if strings.HasPrefix(name, "v.") {
			for _, tensor := range layer {
				weights += tensor.Size()
			}
		}
	}

	switch llm.KV().Architecture() {
	case "mllama":
		kv := func(n string) uint64 {
			if v, ok := llm.KV()["mllama.vision."+n].(uint32); ok {
				return uint64(v)
			}

			return 0
		}

		imageSize := kv("image_size")

		maxNumTiles := kv("max_num_tiles")
		embeddingLength := kv("embedding_length")
		headCount := kv("attention.head_count")

		numPatches := (imageSize / kv("patch_size")) * (imageSize / kv("patch_size"))
		if _, ok := llm.Tensors().GroupLayers()["v"]["class_embd"]; ok {
			numPatches++
		}

		numPaddedPatches := numPatches + 8 - (numPatches%8)%8

		graphSize = 4 * (8 +
			imageSize*imageSize*kv("num_channels")*maxNumTiles +
			embeddingLength*numPatches*maxNumTiles +

Review comment (suggestion):
Doesn't have to be in this PR but we should really add the graph/activation memory requirements if we can

Code after:
				partialOffload,
				4*batch*(1+
					2*embedding+
					embeddingHeadsK*heads+
					context+
					context*heads)+
					4*embeddingHeadsK*context+
					4*context*embeddingHeadsK+
					4*qkvBias.Shape[0],
			)
		}
	}

	return
}

func (llm GGML) VisionGraphSize() (weights, graphSize uint64) {
	if llm.KV().Uint("vision.block_count") == 0 {
		return
	}

	for name, layer := range llm.Tensors().GroupLayers() {
		if name == "v" || strings.HasPrefix(name, "v.") {
			for _, tensor := range layer {
				weights += tensor.Size()
			}
		}
	}

	imageSize := uint64(llm.KV().Uint("vision.image_size"))
	patchSize := uint64(llm.KV().Uint("vision.patch_size"))
	if patchSize == 0 {
		slog.Warn("unknown patch size for vision model")
		return
	}

	numChannels := uint64(llm.KV().Uint("vision.num_channels"))

	numPatches := (imageSize / patchSize) * (imageSize / patchSize)
	if _, ok := llm.Tensors().GroupLayers()["v"]["class_embd"]; ok {
		numPatches++
	}

	headCount := uint64(llm.KV().Uint("vision.attention.head_count"))
	embeddingLength := uint64(llm.KV().Uint("vision.embedding_length"))

	switch llm.KV().Architecture() {
	case "mllama":
		numPaddedPatches := numPatches + 8 - (numPatches%8)%8

		maxNumTiles := uint64(llm.KV().Uint("vision.max_num_tiles"))
Repository: ollama/ollama
File: sample/transforms.go (Go)

Code before:
	"slices"
)

// tokenHeap implements heap.Interface and holds tokens as a min-heap to track k largest elements
type tokenHeap []token

func (h tokenHeap) Len() int           { return len(h) }
func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h tokenHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *tokenHeap) Push(x any) {
	*h = append(*h, x.(token))
}

func (h *tokenHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// temperature applies scaling to the logits
func temperature(ts []token, temp float32) []token {
	temp = max(temp, 1e-7)
	for i := range ts {
		ts[i].value = ts[i].value / temp
	}
	return ts
}

// softmax applies normalization to the logits
func softmax(ts []token) []token {
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Compute exp(x - max)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64(v.value - maxLogit)))
		sum += ts[i].value
	}

	// exp(x - max) / sum(exp(x - max))
	for i := range ts {
		ts[i].value /= sum

Review comment (suggestion):
I don't mind the `max`. For readability we may want to do something like:

```
// avoid division by zero
if t < 1e-7 {
    t = 1e-7
}
```


Code after:
	"slices"
)

// tokenHeap implements heap.Interface and holds tokens as a min-heap to track k largest elements
type tokenHeap []token

func (h tokenHeap) Len() int           { return len(h) }
func (h tokenHeap) Less(i, j int) bool { return h[i].value < h[j].value }
func (h tokenHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *tokenHeap) Push(x any) {
	*h = append(*h, x.(token))
}

func (h *tokenHeap) Pop() any {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

// temperature applies scaling to the logits
func temperature(ts []token, temp float32) []token {
	// Ensure temperature clipping near 0 to avoid numerical instability
	temp = max(temp, 1e-7)
	for i := range ts {
		ts[i].value = ts[i].value / temp
	}
	return ts
}

// softmax applies normalization to the logits
func softmax(ts []token) []token {
	// Find max logit for numerical stability
	maxLogit := float32(math.Inf(-1))
	for _, t := range ts {
		if t.value > maxLogit {
			maxLogit = t.value
		}
	}

	// Compute exp(x - max)
	var sum float32
	for i, v := range ts {
		ts[i].value = float32(math.Exp(float64(v.value - maxLogit)))
		sum += ts[i].value
	}

	// exp(x - max) / sum(exp(x - max))
	for i := range ts {