nx_fnet_train_gate.nx source

1// nx_fnet_train_gate.nx -- GATE for FNET-003: the trainable FNet block (arc convergence). Proves, by RUNNING: 2// G1 f32 FFT correctness: the f32 radix-2 FFT equals the naive O(n^2) f32 DFT (same CORDIC->f32 twiddles) 3// at n=8 -- the f32 transform is a real Fourier transform. 4// G2 FNet-op GRADCHECK through the autograd: x[4,4] -> ta_fnet -> mse(.,target); analytic gradient (the 5// self-adjoint backward, dx = fnet_mix(dL/dy)) vs central finite difference per input cell (h=1/128, 6// rel<1/32 floor 1/64). Proves the sub-quadratic mixer is now a correctly-differentiable tape op. 7// G3 TRAIN through the op: gradient descent on the input to minimize mse(fnet(x), target) -- the loss 8// strictly decreases, so gradients flow end-to-end through the FNet block. 9// 10// Evidence -> knowledge/status/fnet_train.log (FNETTRAINGATE authored=organ ... verdict=GREEN). license_tier: ORIGINAL 11import "nx_autograd_tensor.nx" // ta_fnet + ta_* + transitively nx_fft_f32 (fftf_*, fnet_mix_f32) + fx + f32 tower 12import "nx_syscalls.nx" 13import "nx_gate_verdict.nx" 14 15const FT_LOG: *u8 = "knowledge/status/fnet_train.log" 16 17func ft_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 } 18func ft_wn(fd: i64, v: i64) -> i64 { 19 let bb: *u8 = sys_mmap(28); var m: i64 = v 20 if m < 0 { m = 0 - m; sys_write(fd, "-" as *u8, 1) } 21 let t: *u8 = sys_mmap(28); var k: i64 = 0 22 if m == 0 { t[0] = 48; k = 1 } 23 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 } 24 var i: i64 = 0 25 while i < k { bb[i] = t[k - 1 - i]; i = i + 1 } 26 sys_write(fd, bb, k); return 0 27} 28 29// naive O(n^2) f32 DFT using the same CORDIC->f32 twiddles as the FFT. 30func ft_dft(in_re: *i64, in_im: *i64, n: i64, out_re: *i64, out_im: *i64) -> i64 { 31 var k: i64 = 0 32 while k < n { 33 var sr: i64 = TA_F32_ZERO 34 var si: i64 = TA_F32_ZERO 35 var j: i64 = 0 36 while j < n { 37 let ang: i64 = FX_TWO_PI * ((j * k) % n) / n 38 let wr: i64 = fftf_q16_to_f32(fx_cos(ang)) 39 let wi: i64 = fftf_q16_to_f32(0 - fx_sin(ang)) 40 sr = nx_f32_add(sr, nx_f32_sub(nx_f32_mul(in_re[j], wr), nx_f32_mul(in_im[j], wi))) 41 si = nx_f32_add(si, nx_f32_add(nx_f32_mul(in_re[j], wi), nx_f32_mul(in_im[j], wr))) 42 j = j + 1 43 } 44 out_re[k] = sr; out_im[k] = si 45 k = k + 1 46 } 47 return 0 48} 49 50// ===== Gate 2/3: FNet-op loss graph: loss = mse(ta_fnet(x), target), x and target both [4,4] ===== 51func g_fnet_build(tape: *i64, vals: *i64, st: *i64, x: *i64, target: *i64, xnout: *i64) -> i64 { 52 st[0] = 0; st[1] = 0 53 let xn: i64 = ta_leaf(tape, vals, st, 4, 4, x, 0) 54 let fn: i64 = ta_fnet(tape, vals, st, xn) 55 let tn: i64 = ta_leaf(tape, vals, st, 4, 4, target, 0) 56 let loss: i64 = ta_mse(tape, vals, st, fn, tn) 57 xnout[0] = xn 58 return loss 59} 60func g_fnet_loss(tape: *i64, vals: *i64, st: *i64, x: *i64, target: *i64) -> i64 { 61 let xn: *i64 = (sys_mmap(8)) as *i64 62 let loss: i64 = g_fnet_build(tape, vals, st, x, target, xn) 63 return ta_val(tape, vals, loss, 0) 64} 65func g_fnet_grads(tape: *i64, vals: *i64, grads: *i64, st: *i64, x: *i64, target: *i64, dx: *i64) -> i64 { 66 let xn: *i64 = (sys_mmap(8)) as *i64 67 let loss: i64 = g_fnet_build(tape, vals, st, x, target, xn) 68 ta_backward(tape, vals, grads, st[0], loss) 69 var i: i64 = 0 70 while i < 16 { dx[i] = ta_grad(tape, grads, xn[0], i); i = i + 1 } 71 return 0 72} 73 74func main() -> i64 { 75 var ok: i64 = 1 76 let tape: *i64 = (sys_mmap(256 * 7 * 8)) as *i64 77 let vals: *i64 = (sys_mmap(4096 * 8)) as *i64 78 let grads: *i64 = (sys_mmap(4096 * 8)) as *i64 79 let st: *i64 = (sys_mmap(2 * 8)) as *i64 80 81 // ---------- G1: f32 FFT == naive f32 DFT at n=8 ---------- 82 let n1: i64 = 8 83 let ar: *i64 = (sys_mmap(n1 * 8)) as *i64 84 let ai: *i64 = (sys_mmap(n1 * 8)) as *i64 85 let br: *i64 = (sys_mmap(n1 * 8)) as *i64 86 let bi: *i64 = (sys_mmap(n1 * 8)) as *i64 87 var i: i64 = 0 88 while i < n1 { 89 let v: i64 = ta_constf((i % 5) - 2, 1) 90 ar[i] = v; ai[i] = TA_F32_ZERO; br[i] = v; bi[i] = TA_F32_ZERO 91 i = i + 1 92 } 93 let twr: *i64 = (sys_mmap(n1 * 8)) as *i64 94 let twi: *i64 = (sys_mmap(n1 * 8)) as *i64 95 fftf_twiddles(n1, twr, twi) 96 fftf_fwd(ar, ai, n1, twr, twi) 97 let dor: *i64 = (sys_mmap(n1 * 8)) as *i64 98 let doi: *i64 = (sys_mmap(n1 * 8)) as *i64 99 ft_dft(br, bi, n1, dor, doi) 100 var g1max: i64 = TA_F32_ZERO 101 i = 0 102 while i < n1 { 103 let d1: i64 = nx_f32_abs(nx_f32_sub(ar[i], dor[i])) 104 let d2: i64 = nx_f32_abs(nx_f32_sub(ai[i], doi[i])) 105 if nx_f32_gt(d1, g1max) == 1 { g1max = d1 } 106 if nx_f32_gt(d2, g1max) == 1 { g1max = d2 } 107 i = i + 1 108 } 109 let g1milli: i64 = ta_f32_to_milli(g1max) 110 if g1milli > 50 { ok = 0 } // FFT vs DFT agree to < 0.05 (f32 rounding-order only) 111 112 // ---------- G2: FNet-op gradcheck through the autograd ---------- 113 let x: *i64 = (sys_mmap(16 * 8)) as *i64 114 let target: *i64 = (sys_mmap(16 * 8)) as *i64 115 i = 0 116 while i < 16 { 117 x[i] = ta_constf((i % 5) - 2, 2) // -1 .. 1 118 target[i] = ta_constf((i % 7) - 3, 4) // -0.75 .. 0.75 119 i = i + 1 120 } 121 let ana: *i64 = (sys_mmap(16 * 8)) as *i64 122 g_fnet_grads(tape, vals, grads, st, x, target, ana) 123 let h: i64 = ta_constf(1, 128) 124 let flo: i64 = ta_constf(1, 64) 125 let tol: i64 = ta_constf(1, 32) 126 let xp: *i64 = (sys_mmap(16 * 8)) as *i64 127 let xm: *i64 = (sys_mmap(16 * 8)) as *i64 128 var g2pass: i64 = 1 129 var g2worst: i64 = 0 130 var c: i64 = 0 131 while c < 16 { 132 var j: i64 = 0 133 while j < 16 { xp[j] = x[j]; xm[j] = x[j]; j = j + 1 } 134 xp[c] = nx_f32_add(x[c], h) 135 xm[c] = nx_f32_sub(x[c], h) 136 let lp: i64 = g_fnet_loss(tape, vals, st, xp, target) 137 let lm: i64 = g_fnet_loss(tape, vals, st, xm, target) 138 let fd: i64 = nx_f32_div(nx_f32_sub(lp, lm), nx_f32_add(h, h)) 139 let num: i64 = nx_f32_abs(nx_f32_sub(fd, ana[c])) 140 var den: i64 = nx_f32_abs(ana[c]) 141 if nx_f32_lt(den, flo) == 1 { den = flo } 142 if nx_f32_lt(num, nx_f32_mul(tol, den)) != 1 { g2pass = 0 } 143 let nm: i64 = ta_f32_to_milli(num) 144 if nm > g2worst { g2worst = nm } 145 c = c + 1 146 } 147 if g2pass != 1 { ok = 0 } 148 149 // ---------- G3: train the input through the FNet op (loss must strictly decrease) ---------- 150 let xt: *i64 = (sys_mmap(16 * 8)) as *i64 151 i = 0 152 while i < 16 { xt[i] = TA_F32_ZERO; i = i + 1 } 153 let lr: i64 = ta_constf(1, 64) 154 let dxt: *i64 = (sys_mmap(16 * 8)) as *i64 155 var lfirst: i64 = TA_F32_ZERO 156 var llast: i64 = TA_F32_ZERO 157 var ep: i64 = 0 158 while ep < 1000 { 159 if ep == 0 { lfirst = g_fnet_loss(tape, vals, st, xt, target) } 160 g_fnet_grads(tape, vals, grads, st, xt, target, dxt) 161 i = 0 162 while i < 16 { xt[i] = nx_f32_sub(xt[i], nx_f32_mul(lr, dxt[i])); i = i + 1 } 163 ep = ep + 1 164 } 165 llast = g_fnet_loss(tape, vals, st, xt, target) 166 var g3pass: i64 = 1 167 if nx_f32_lt(llast, lfirst) != 1 { g3pass = 0 } // loss strictly decreased 168 if g3pass != 1 { ok = 0 } 169 170 // ---------- emit ---------- 171 var fdi: i64 = 1 172 while fdi >= 0 { 173 var out: i64 = 1 174 if fdi == 0 { out = sys_openat_append(FT_LOG, 420) } 175 if out >= 0 { 176 ft_w(out, "FNETTRAINGATE authored=organ block=trainable-fnet-f32" as *u8) 177 ft_w(out, " | G1_fft_vs_dft_milli=" as *u8); ft_wn(out, g1milli) 178 ft_w(out, " | G2_gradcheck_pass=" as *u8); ft_wn(out, g2pass); ft_w(out, " worst_|fd-analytic|_milli=" as *u8); ft_wn(out, g2worst) 179 ft_w(out, " | G3_train_pass=" as *u8); ft_wn(out, g3pass) 180 ft_w(out, " loss_first_milli=" as *u8); ft_wn(out, ta_f32_to_milli(lfirst)) 181 ft_w(out, " loss_last_milli=" as *u8); ft_wn(out, ta_f32_to_milli(llast)) 182 if ok == 1 { ft_w(out, " verdict=GREEN\n" as *u8) } else { ft_w(out, " verdict=RED\n" as *u8) } 183 if fdi == 0 { sys_close(out) } 184 } 185 fdi = fdi - 1 186 } 187 188 // MIGRATED onto nx_gate_verdict by nx_gate_dry_apply (D001, minimal form): every check 189 // row above is untouched, so the PASS/FAIL vector cannot change; only the hand-rolled 190 // verdict emission is replaced by the ONE shared base class. Proven by nx_gate_migrate verify. 191 let ctr__dry: *i64 = gv_ctr() 192 ctr__dry[0] = ok 193 ctr__dry[1] = 1 194 let rc__dry: i64 = gv_verdict("FNET-TRAIN-GATE" as *u8, ctr__dry, "teeth unchanged; verdict emission migrated onto the shared base class" as *u8) 195 sys_exit(rc__dry) 196 return rc__dry 197}