code wiki / _hdl_build / nx_train_r3_gate.nx
nx_train_r3_gate.nx source
↩ module page · 292 lines · 13865 B
1// nx_train_r3_gate.nx -- GATE for TRAIN-R3 (T7): softmax + cross-entropy + nonlinear training. Proves, by RUNNING:
2// A GRADCHECK the new ops: softmax-cross-entropy (dlogits = softmax - onehot) AND standalone softmax
3// (Jacobian-vector product), analytic vs central finite difference (h=1/128, rel<1/32 floor 1/64).
4// B TRAIN A NONLINEAR CLASSIFIER -- XOR: a 2->8(relu)->2 MLP with a softmax-CE head and a DETERMINISTIC
5// symmetry-breaking init, trained by AdamW. XOR is the canonical proof that a LINEAR model CANNOT solve
6// it -- reaching 4/4 correct demonstrates real nonlinear learning. Assert accuracy==4 AND loss decreased.
7// C BIT-EXACT: train twice from the same init -> identical final bits for all 42 parameters.
8//
9// Evidence -> knowledge/status/train_r3.log (TRAINR3GATE authored=organ ... verdict=GREEN). license_tier: ORIGINAL
10import "nx_autograd_tensor.nx" // ta_* (now incl softmax/softce/det_init) + transitively the f32 tower / syscalls
11import "nx_syscalls.nx"
12import "nx_gate_verdict.nx"
13
14const T3_LOG: *u8 = "knowledge/status/train_r3.log"
15
16func t3_w(fd: i64, s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(fd, s, n); return 0 }
17func t3_wn(fd: i64, v: i64) -> i64 {
18 let bb: *u8 = sys_mmap(28); var m: i64 = v
19 if m < 0 { m = 0 - m; sys_write(fd, "-" as *u8, 1) }
20 let t: *u8 = sys_mmap(28); var k: i64 = 0
21 if m == 0 { t[0] = 48; k = 1 }
22 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
23 var i: i64 = 0
24 while i < k { bb[i] = t[k - 1 - i]; i = i + 1 }
25 sys_write(fd, bb, k); return 0
26}
27
28// ===== Gate A: softmax-CE gradcheck =====
29func g_softce_loss(tape: *i64, vals: *i64, st: *i64, lg: *i64, tg: *i64) -> i64 {
30 st[0] = 0; st[1] = 0
31 let ln: i64 = ta_leaf(tape, vals, st, 3, 1, lg, 0)
32 let tn: i64 = ta_leaf(tape, vals, st, 3, 1, tg, 0)
33 let loss: i64 = ta_softce(tape, vals, st, ln, tn)
34 return ta_val(tape, vals, loss, 0)
35}
36func g_softce_grads(tape: *i64, vals: *i64, grads: *i64, st: *i64, lg: *i64, tg: *i64, gout: *i64) -> i64 {
37 st[0] = 0; st[1] = 0
38 let ln: i64 = ta_leaf(tape, vals, st, 3, 1, lg, 0)
39 let tn: i64 = ta_leaf(tape, vals, st, 3, 1, tg, 0)
40 let loss: i64 = ta_softce(tape, vals, st, ln, tn)
41 ta_backward(tape, vals, grads, st[0], loss)
42 gout[0] = ta_grad(tape, grads, ln, 0); gout[1] = ta_grad(tape, grads, ln, 1); gout[2] = ta_grad(tape, grads, ln, 2)
43 return 0
44}
45// ===== Gate A: standalone softmax gradcheck (loss = mse(softmax(x), tg)) =====
46func g_smax_loss(tape: *i64, vals: *i64, st: *i64, x: *i64, tg: *i64) -> i64 {
47 st[0] = 0; st[1] = 0
48 let xn: i64 = ta_leaf(tape, vals, st, 3, 1, x, 0)
49 let sm: i64 = ta_softmax(tape, vals, st, xn)
50 let tn: i64 = ta_leaf(tape, vals, st, 3, 1, tg, 0)
51 let loss: i64 = ta_mse(tape, vals, st, sm, tn)
52 return ta_val(tape, vals, loss, 0)
53}
54func g_smax_grads(tape: *i64, vals: *i64, grads: *i64, st: *i64, x: *i64, tg: *i64, gout: *i64) -> i64 {
55 st[0] = 0; st[1] = 0
56 let xn: i64 = ta_leaf(tape, vals, st, 3, 1, x, 0)
57 let sm: i64 = ta_softmax(tape, vals, st, xn)
58 let tn: i64 = ta_leaf(tape, vals, st, 3, 1, tg, 0)
59 let loss: i64 = ta_mse(tape, vals, st, sm, tn)
60 ta_backward(tape, vals, grads, st[0], loss)
61 gout[0] = ta_grad(tape, grads, xn, 0); gout[1] = ta_grad(tape, grads, xn, 1); gout[2] = ta_grad(tape, grads, xn, 2)
62 return 0
63}
64
65// generic 3-vector gradcheck: returns 1 if all pass, writes worst |fd-analytic| milli into *worst. kind 0=softce,1=softmax.
66func g_gc3(tape: *i64, vals: *i64, grads: *i64, st: *i64, base: *i64, tg: *i64, kind: i64, worst: *i64) -> i64 {
67 let ana: *i64 = (sys_mmap(3 * 8)) as *i64
68 if kind == 0 { g_softce_grads(tape, vals, grads, st, base, tg, ana) } else { g_smax_grads(tape, vals, grads, st, base, tg, ana) }
69 let h: i64 = ta_constf(1, 128)
70 let flo: i64 = ta_constf(1, 64)
71 let tol: i64 = ta_constf(1, 32)
72 let pp: *i64 = (sys_mmap(3 * 8)) as *i64
73 let pm: *i64 = (sys_mmap(3 * 8)) as *i64
74 var pass: i64 = 1
75 var wm: i64 = 0
76 var pi: i64 = 0
77 while pi < 3 {
78 var j: i64 = 0
79 while j < 3 { pp[j] = base[j]; pm[j] = base[j]; j = j + 1 }
80 pp[pi] = nx_f32_add(base[pi], h)
81 pm[pi] = nx_f32_sub(base[pi], h)
82 var lp: i64 = 0
83 var lm: i64 = 0
84 if kind == 0 { lp = g_softce_loss(tape, vals, st, pp, tg); lm = g_softce_loss(tape, vals, st, pm, tg) }
85 else { lp = g_smax_loss(tape, vals, st, pp, tg); lm = g_smax_loss(tape, vals, st, pm, tg) }
86 let fd: i64 = nx_f32_div(nx_f32_sub(lp, lm), nx_f32_add(h, h))
87 let num: i64 = nx_f32_abs(nx_f32_sub(fd, ana[pi]))
88 var den: i64 = nx_f32_abs(ana[pi])
89 if nx_f32_lt(den, flo) == 1 { den = flo }
90 if nx_f32_lt(num, nx_f32_mul(tol, den)) != 1 { pass = 0 }
91 let nm: i64 = ta_f32_to_milli(num)
92 if nm > wm { wm = nm }
93 pi = pi + 1
94 }
95 *worst = wm
96 return pass
97}
98
99// ===== Gate B: XOR MLP (2 -> 8 relu -> 2, softmax-CE). params p[42]: W1[0..15] b1[16..23] W2[24..39] b2[40..41] =====
100func g_xor_build(tape: *i64, vals: *i64, st: *i64, p: *i64, xs: *i64, ts: *i64, c4: *i64, wb: *i64) -> i64 {
101 st[0] = 0; st[1] = 0
102 let W1: i64 = ta_leaf(tape, vals, st, 8, 2, p, 0)
103 let b1: i64 = ta_leaf(tape, vals, st, 8, 1, p, 16)
104 let W2: i64 = ta_leaf(tape, vals, st, 2, 8, p, 24)
105 let b2: i64 = ta_leaf(tape, vals, st, 2, 1, p, 40)
106 var sumn: i64 = 0 - 1
107 var s: i64 = 0
108 while s < 4 {
109 let x: i64 = ta_leaf(tape, vals, st, 2, 1, xs, s * 2)
110 let h: i64 = ta_relu(tape, vals, st, ta_vadd(tape, vals, st, ta_matvec(tape, vals, st, W1, x), b1))
111 let lo: i64 = ta_vadd(tape, vals, st, ta_matvec(tape, vals, st, W2, h), b2)
112 let tt: i64 = ta_leaf(tape, vals, st, 2, 1, ts, s * 2)
113 let ls: i64 = ta_softce(tape, vals, st, lo, tt)
114 if s == 0 { sumn = ls } else { sumn = ta_vadd(tape, vals, st, sumn, ls) }
115 s = s + 1
116 }
117 let inv4: i64 = ta_leaf(tape, vals, st, 1, 1, c4, 0)
118 let loss: i64 = ta_matvec(tape, vals, st, inv4, sumn)
119 wb[0] = W1; wb[1] = b1; wb[2] = W2; wb[3] = b2
120 return loss
121}
122func g_read42(tape: *i64, grads: *i64, wb: *i64, g: *i64) -> i64 {
123 var i: i64 = 0
124 while i < 16 { g[i] = ta_grad(tape, grads, wb[0], i); i = i + 1 }
125 i = 0
126 while i < 8 { g[16 + i] = ta_grad(tape, grads, wb[1], i); i = i + 1 }
127 i = 0
128 while i < 16 { g[24 + i] = ta_grad(tape, grads, wb[2], i); i = i + 1 }
129 g[40] = ta_grad(tape, grads, wb[3], 0); g[41] = ta_grad(tape, grads, wb[3], 1)
130 return 0
131}
132// argmax of the 2 output logits for sample s
133func g_xor_predict(tape: *i64, vals: *i64, st: *i64, p: *i64, xs: *i64, s: i64) -> i64 {
134 st[0] = 0; st[1] = 0
135 let W1: i64 = ta_leaf(tape, vals, st, 8, 2, p, 0)
136 let b1: i64 = ta_leaf(tape, vals, st, 8, 1, p, 16)
137 let W2: i64 = ta_leaf(tape, vals, st, 2, 8, p, 24)
138 let b2: i64 = ta_leaf(tape, vals, st, 2, 1, p, 40)
139 let x: i64 = ta_leaf(tape, vals, st, 2, 1, xs, s * 2)
140 let h: i64 = ta_relu(tape, vals, st, ta_vadd(tape, vals, st, ta_matvec(tape, vals, st, W1, x), b1))
141 let lo: i64 = ta_vadd(tape, vals, st, ta_matvec(tape, vals, st, W2, h), b2)
142 if nx_f32_gt(ta_val(tape, vals, lo, 1), ta_val(tape, vals, lo, 0)) == 1 { return 1 }
143 return 0
144}
145// init the 42 params: W1 (seed 3) + W2 (seed 7) spread, biases zero.
146func g_xor_init(p: *i64) -> i64 {
147 var i: i64 = 0
148 while i < 42 { p[i] = TA_F32_ZERO; i = i + 1 }
149 i = 0
150 while i < 16 { p[i] = ta_constf(((i * 3 + 1) % 11) - 5, 8); i = i + 1 }
151 i = 0
152 while i < 16 { p[24 + i] = ta_constf(((i * 7 + 1) % 11) - 5, 8); i = i + 1 }
153 return 0
154}
155// AdamW training of the XOR MLP; writes final params + first/last loss.
156func g_xor_train(tape: *i64, vals: *i64, grads: *i64, st: *i64, xs: *i64, ts: *i64, epochs: i64, pout: *i64, lf: *i64, ll: *i64) -> i64 {
157 let p: *i64 = (sys_mmap(42 * 8)) as *i64
158 let mm: *i64 = (sys_mmap(42 * 8)) as *i64
159 let vv: *i64 = (sys_mmap(42 * 8)) as *i64
160 g_xor_init(p)
161 var i: i64 = 0
162 while i < 42 { mm[i] = TA_F32_ZERO; vv[i] = TA_F32_ZERO; i = i + 1 }
163 let beta1: i64 = ta_constf(9, 10)
164 let beta2: i64 = ta_constf(999, 1000)
165 let om1: i64 = ta_constf(1, 10)
166 let om2: i64 = ta_constf(1, 1000)
167 let lr: i64 = ta_constf(1, 20)
168 let eps: i64 = ta_constf(1, 100000000)
169 var b1t: i64 = TA_F32_ONE
170 var b2t: i64 = TA_F32_ONE
171 let c4: *i64 = (sys_mmap(8)) as *i64; c4[0] = ta_constf(1, 4)
172 let wb: *i64 = (sys_mmap(4 * 8)) as *i64
173 let g: *i64 = (sys_mmap(42 * 8)) as *i64
174 var ep: i64 = 0
175 while ep < epochs {
176 let loss: i64 = g_xor_build(tape, vals, st, p, xs, ts, c4, wb)
177 ta_backward(tape, vals, grads, st[0], loss)
178 if ep == 0 { *lf = ta_val(tape, vals, loss, 0) }
179 *ll = ta_val(tape, vals, loss, 0)
180 g_read42(tape, grads, wb, g)
181 b1t = nx_f32_mul(b1t, beta1)
182 b2t = nx_f32_mul(b2t, beta2)
183 let c1c: i64 = nx_f32_sub(TA_F32_ONE, b1t)
184 let c2c: i64 = nx_f32_sub(TA_F32_ONE, b2t)
185 i = 0
186 while i < 42 {
187 let gi: i64 = g[i]
188 mm[i] = nx_f32_add(nx_f32_mul(beta1, mm[i]), nx_f32_mul(om1, gi))
189 vv[i] = nx_f32_add(nx_f32_mul(beta2, vv[i]), nx_f32_mul(om2, nx_f32_mul(gi, gi)))
190 let mhat: i64 = nx_f32_div(mm[i], c1c)
191 let vhat: i64 = nx_f32_div(vv[i], c2c)
192 p[i] = nx_f32_sub(p[i], nx_f32_mul(lr, nx_f32_div(mhat, nx_f32_add(nx_f32_sqrt(vhat), eps))))
193 i = i + 1
194 }
195 ep = ep + 1
196 }
197 i = 0
198 while i < 42 { pout[i] = p[i]; i = i + 1 }
199 return 0
200}
201
202func main() -> i64 {
203 var ok: i64 = 1
204 let tape: *i64 = (sys_mmap(1024 * 7 * 8)) as *i64
205 let vals: *i64 = (sys_mmap(4096 * 8)) as *i64
206 let grads: *i64 = (sys_mmap(4096 * 8)) as *i64
207 let st: *i64 = (sys_mmap(2 * 8)) as *i64
208
209 // ----- Gate A: gradcheck softmax-CE + standalone softmax -----
210 let lg: *i64 = (sys_mmap(3 * 8)) as *i64
211 lg[0] = ta_constf(1, 2); lg[1] = ta_constf(0 - 3, 10); lg[2] = ta_constf(4, 5)
212 let oneh: *i64 = (sys_mmap(3 * 8)) as *i64
213 oneh[0] = TA_F32_ZERO; oneh[1] = TA_F32_ONE; oneh[2] = TA_F32_ZERO
214 let wce: *i64 = (sys_mmap(8)) as *i64
215 let ceGc: i64 = g_gc3(tape, vals, grads, st, lg, oneh, 0, wce)
216 if ceGc != 1 { ok = 0 }
217
218 let xg: *i64 = (sys_mmap(3 * 8)) as *i64
219 xg[0] = ta_constf(2, 5); xg[1] = ta_constf(0 - 1, 5); xg[2] = ta_constf(7, 10)
220 let dist: *i64 = (sys_mmap(3 * 8)) as *i64
221 dist[0] = ta_constf(3, 10); dist[1] = ta_constf(2, 10); dist[2] = ta_constf(5, 10)
222 let wsm: *i64 = (sys_mmap(8)) as *i64
223 let smGc: i64 = g_gc3(tape, vals, grads, st, xg, dist, 1, wsm)
224 if smGc != 1 { ok = 0 }
225
226 // ----- Gate B: train XOR -----
227 let xs: *i64 = (sys_mmap(8 * 8)) as *i64
228 xs[0] = TA_F32_ZERO; xs[1] = TA_F32_ZERO
229 xs[2] = TA_F32_ZERO; xs[3] = TA_F32_ONE
230 xs[4] = TA_F32_ONE; xs[5] = TA_F32_ZERO
231 xs[6] = TA_F32_ONE; xs[7] = TA_F32_ONE
232 let ts: *i64 = (sys_mmap(8 * 8)) as *i64 // onehot targets per XOR label 0,1,1,0
233 ts[0] = TA_F32_ONE; ts[1] = TA_F32_ZERO
234 ts[2] = TA_F32_ZERO; ts[3] = TA_F32_ONE
235 ts[4] = TA_F32_ZERO; ts[5] = TA_F32_ONE
236 ts[6] = TA_F32_ONE; ts[7] = TA_F32_ZERO
237 let labels: *i64 = (sys_mmap(4 * 8)) as *i64
238 labels[0] = 0; labels[1] = 1; labels[2] = 1; labels[3] = 0
239 let pB: *i64 = (sys_mmap(42 * 8)) as *i64
240 let lfB: *i64 = (sys_mmap(8)) as *i64
241 let llB: *i64 = (sys_mmap(8)) as *i64
242 g_xor_train(tape, vals, grads, st, xs, ts, 2500, pB, lfB, llB)
243 var acc: i64 = 0
244 var s: i64 = 0
245 while s < 4 {
246 if g_xor_predict(tape, vals, st, pB, xs, s) == labels[s] { acc = acc + 1 }
247 s = s + 1
248 }
249 var learnsB: i64 = 1
250 if acc != 4 { learnsB = 0 }
251 if nx_f32_lt(*llB, *lfB) != 1 { learnsB = 0 }
252 if learnsB != 1 { ok = 0 }
253
254 // ----- Gate C: bit-exact -----
255 let pC: *i64 = (sys_mmap(42 * 8)) as *i64
256 let lfC: *i64 = (sys_mmap(8)) as *i64
257 let llC: *i64 = (sys_mmap(8)) as *i64
258 g_xor_train(tape, vals, grads, st, xs, ts, 2500, pC, lfC, llC)
259 var reproC: i64 = 1
260 var i: i64 = 0
261 while i < 42 { if pC[i] != pB[i] { reproC = 0 } i = i + 1 }
262 if reproC != 1 { ok = 0 }
263
264 // ----- emit -----
265 var fdi: i64 = 1
266 while fdi >= 0 {
267 var out: i64 = 1
268 if fdi == 0 { out = sys_openat_append(T3_LOG, 420) }
269 if out >= 0 {
270 t3_w(out, "TRAINR3GATE authored=organ engine=tensor-autograd softmax+crossentropy+nonlinear" as *u8)
271 t3_w(out, " | A_softce_gradcheck=" as *u8); t3_wn(out, ceGc); t3_w(out, " worst_milli=" as *u8); t3_wn(out, wce[0])
272 t3_w(out, " A_softmax_gradcheck=" as *u8); t3_wn(out, smGc); t3_w(out, " worst_milli=" as *u8); t3_wn(out, wsm[0])
273 t3_w(out, " | B_XOR_accuracy=" as *u8); t3_wn(out, acc); t3_w(out, "/4" as *u8)
274 t3_w(out, " loss_first_milli=" as *u8); t3_wn(out, ta_f32_to_milli(*lfB))
275 t3_w(out, " loss_last_milli=" as *u8); t3_wn(out, ta_f32_to_milli(*llB))
276 t3_w(out, " | C_bitexact=" as *u8); t3_wn(out, reproC)
277 if ok == 1 { t3_w(out, " verdict=GREEN\n" as *u8) } else { t3_w(out, " verdict=RED\n" as *u8) }
278 if fdi == 0 { sys_close(out) }
279 }
280 fdi = fdi - 1
281 }
282
283 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check
284 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled
285 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify.
286 let ctr__dry: *i64 = gv_ctr()
287 ctr__dry[0] = ok
288 ctr__dry[1] = 1
289 let rc__dry: i64 = gv_verdict("TRAIN-R3-GATE" as *u8, ctr__dry, "teeth unchanged; verdict emission migrated onto the shared base class" as *u8)
290 sys_exit(rc__dry)
291 return rc__dry
292}