js/src/runtime/softmax.ts

/**
 * Numerically stable softmax + argmax-with-confidence.
 *
 * Pure JS — no ORT, no DOM. The classifier ONNX emits raw logits
 * because softmax-inside-graph would erase the per-class info we
 * use to surface ``confidence`` to downstream callers.
 */

/** Argmax + softmax probability of the winning class. Numerically
 * stable (subtract max before exp). */
export function argmaxWithProb(
  logits: Float32Array | number[],
): { idx: number; prob: number } {
  let maxLogit = -Infinity;
  let argmax = 0;
  for (let i = 0; i < logits.length; i++) {
    if (logits[i] > maxLogit) {
      maxLogit = logits[i];
      argmax = i;
    }
  }
  let z = 0;
  for (let i = 0; i < logits.length; i++) {
    z += Math.exp(logits[i] - maxLogit);
  }
  return { idx: argmax, prob: 1 / z };
}