INNER CODE UNIT · Go
scale
takara-ai/go-attention · attention/attention.go:191
scale = 1.0 / math.Sqrt(float64(d_k))
} // If d_k is 0, scale remains 1.0, dot product will likely be 0 unless vectors are empty.
for i, key := range keys {
// Ensure consistent key dimensions within the loop
if len(key) != d_k {
return nil, nil, fmt.Errorf("key dimension mismatch at index %d: expected %d, got %d", i, d_k, len(key))
}
score, err := DotProduct(query, key) // DotProduct already checks len(query) == len(key)
if err != nil {
// This error should theoretically not happen if the outer checks pass, but handle defensively.
return nil, nil, fmt.Errorf("error computing dot product for key %d: %w", i, err)
}
// Scale by sqrt(d_k) for better gradient flow
scores[i] = score * scale // Use pre-calculated scale
}
// --- SSMax Modification (s=1) ---