chickadee » nanograd » softmax

Softmax normalization with numerical stability and batch support.

Input shapes:

• 1D: (n_classes,) - standard softmax
• 2D: (batch_size, n_classes) - softmax along axis (default: -1 for last axis)

; Single sample
(define logits (make-tensor32 (f32vector 1.0 2.0 3.0) '(3)))
(define probs (softmax logits))  ; Sums to 1.0

; Batch of samples
(define batch-logits (make-tensor32 (make-f32vector 60) '(20 3)))
(define batch-probs (softmax batch-logits axis: -1))  ; Each row sums to 1.0

Gradient: dL/dx = softmax(x) ⊙ (dL/dy - Σ(dL/dy ⊙ softmax(x)))