softmax.ts

/**
 * Softmax: Convert logits to probabilities
 *
 * Formula: prob[i] = exp(logit[i]) / sum(exp(logit[j]) for all j)
 *
 * Uses max subtraction for numerical stability:
 * exp(x - max) prevents overflow for large logits
 *
 * @param logits Logits array
 * @returns Probability distribution (sums to 1.0)
 */
export function softmax(logits: Float32Array): Float32Array {
  const probs = new Float32Array(logits.length);

  // Find max for numerical stability
  let maxLogit = -Infinity;
  for (let i = 0; i < logits.length; i++) {
    if (logits[i] > maxLogit) maxLogit = logits[i];
  }

  // Compute exp and sum
  let sum = 0;
  for (let i = 0; i < logits.length; i++) {
    probs[i] = Math.exp(logits[i] - maxLogit);
    sum += probs[i];
  }

  // Normalize
  for (let i = 0; i < logits.length; i++) {
    probs[i] /= sum;
  }

  return probs;
}