nx_f32_softmax.nx source
↩ module page · 87 lines · 2474 B
1// nx_f32_softmax.nx -- bits-up f32 numerically-stable softmax.
2//
3// L7 / L8 composition brick. Composes L4 mul/add/sub/div + L6 exp.
4//
5// Algorithm (numerically stable via max-subtraction):
6// m = max_i x[i]
7// shifted_i = x[i] - m
8// exps_i = exp(shifted_i) (each in (0, 1])
9// sum_exp = sum_i exps_i
10// out[i] = exps_i / sum_exp
11//
12// Max-subtraction guards against exp overflow without changing the
13// final probabilities (numerator+denominator scale by the same
14// 1/exp(m) factor, which cancels).
15//
16// References absorbed clean-room:
17// Bridle 1990 (original softmax)
18// Standard logsumexp / max-subtract trick (numerical canon)
19//
20// genealogy_id: bridle_1990_softmax + logsumexp_max_subtract
21// lineage_id: substrate_f32_softmax_v1_stable
22
23import "nx_syscalls.nx"
24import "nx_tier.nx"
25import "nx_f32.nx"
26import "nx_f32_div.nx"
27import "nx_f32_exp.nx"
28
29const NX_F32_SM_OK: nx_int = 0
30const NX_F32_SM_ERR_BAD_DIM: nx_int = 1
31const NX_F32_SM_N_VERDICTS: nx_int = 2
32
33func nx_f32_sm_verdict_is_valid(v: nx_int) -> nx_int {
34 if v < 0 { return 0 }
35 if v >= NX_F32_SM_N_VERDICTS { return 0 }
36 return 1
37}
38
39// Returns 1 iff a > b in f32 (NaN-naive; we don't expect NaN inputs).
40// Compose-only -- subtract and check sign + non-zero.
41
42func _f32_gt(a: i64, b: i64) -> nx_int {
43 let diff: i64 = nx_f32_sub(a, b)
44 let cls: nx_int = nx_f32_classify(diff)
45 if cls == NX_F32_CLS_ZERO { return 0 }
46 if cls == NX_F32_CLS_NAN { return 0 }
47 if nx_f32_sign(diff) == 1 { return 0 }
48 return 1
49}
50
51// Stable softmax over n f32 values. In-place permitted (out == x).
52//
53// Args:
54// x n f32 input raw-bit values
55// n positive count
56// out n f32 output raw-bit slots
57
58func nx_f32_softmax(x: *i64, n: nx_int, out: *i64) -> nx_int {
59 if n <= 0 { return NX_F32_SM_ERR_BAD_DIM }
60
61 // Pass 1: find max.
62 var m: i64 = x[0]
63 var i: nx_int = 1
64 while i < n {
65 if _f32_gt(x[i], m) == 1 { m = x[i] }
66 i = i + 1
67 }
68
69 // Pass 2: exp(x[i] - m) into out; accumulate sum.
70 var sum: i64 = 0
71 var j: nx_int = 0
72 while j < n {
73 let shifted: i64 = nx_f32_sub(x[j], m)
74 let e: i64 = nx_f32_exp(shifted)
75 out[j] = e
76 sum = nx_f32_add(sum, e)
77 j = j + 1
78 }
79
80 // Pass 3: divide by sum.
81 var k: nx_int = 0
82 while k < n {
83 out[k] = nx_f32_div(out[k], sum)
84 k = k + 1
85 }
86 return NX_F32_SM_OK
87}