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" 13 14const FT_LOG: *u8 = "knowledge/status/fnet_train.log" 15 16func 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 } 17func ft_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// naive O(n^2) f32 DFT using the same CORDIC->f32 twiddles as the FFT. 29func ft_dft(in_re: *i64, in_im: *i64, n: i64, out_re: *i64, out_im: *i64) -> i64 { 30 var k: i64 = 0 31 while k < n { 32 var sr: i64 = TA_F32_ZERO 33 var si: i64 = TA_F32_ZERO 34 var j: i64 = 0 35 while j < n { 36 let ang: i64 = FX_TWO_PI * ((j * k) % n) / n 37 let wr: i64 = fftf_q16_to_f32(fx_cos(ang)) 38 let wi: i64 = fftf_q16_to_f32(0 - fx_sin(ang)) 39 sr = nx_f32_add(sr, nx_f32_sub(nx_f32_mul(in_re[j], wr), nx_f32_mul(in_im[j], wi))) 40 si = nx_f32_add(si, nx_f32_add(nx_f32_mul(in_re[j], wi), nx_f32_mul(in_im[j], wr))) 41 j = j + 1 42 } 43 out_re[k] = sr; out_im[k] = si 44 k = k + 1 45 } 46 return 0 47} 48 49// ===== Gate 2/3: FNet-op loss graph: loss = mse(ta_fnet(x), target), x and target both [4,4] ===== 50func g_fnet_build(tape: *i64, vals: *i64, st: *i64, x: *i64, target: *i64, xnout: *i64) -> i64 { 51 st[0] = 0; st[1] = 0 52 let xn: i64 = ta_leaf(tape, vals, st, 4, 4, x, 0) 53 let fn: i64 = ta_fnet(tape, vals, st, xn) 54 let tn: i64 = ta_leaf(tape, vals, st, 4, 4, target, 0) 55 let loss: i64 = ta_mse(tape, vals, st, fn, tn) 56 xnout[0] = xn 57 return loss 58} 59func g_fnet_loss(tape: *i64, vals: *i64, st: *i64, x: *i64, target: *i64) -> i64 { 60 let xn: *i64 = (sys_mmap(8)) as *i64 61 let loss: i64 = g_fnet_build(tape, vals, st, x, target, xn) 62 return ta_val(tape, vals, loss, 0) 63} 64func g_fnet_grads(tape: *i64, vals: *i64, grads: *i64, st: *i64, x: *i64, target: *i64, dx: *i64) -> i64 { 65 let xn: *i64 = (sys_mmap(8)) as *i64 66 let loss: i64 = g_fnet_build(tape, vals, st, x, target, xn) 67 ta_backward(tape, vals, grads, st[0], loss) 68 var i: i64 = 0 69 while i < 16 { dx[i] = ta_grad(tape, grads, xn[0], i); i = i + 1 } 70 return 0 71} 72 73func main() -> i64 { 74 var ok: i64 = 1 75 let tape: *i64 = (sys_mmap(256 * 7 * 8)) as *i64 76 let vals: *i64 = (sys_mmap(4096 * 8)) as *i64 77 let grads: *i64 = (sys_mmap(4096 * 8)) as *i64 78 let st: *i64 = (sys_mmap(2 * 8)) as *i64 79 80 // ---------- G1: f32 FFT == naive f32 DFT at n=8 ---------- 81 let n1: i64 = 8 82 let ar: *i64 = (sys_mmap(n1 * 8)) as *i64 83 let ai: *i64 = (sys_mmap(n1 * 8)) as *i64 84 let br: *i64 = (sys_mmap(n1 * 8)) as *i64 85 let bi: *i64 = (sys_mmap(n1 * 8)) as *i64 86 var i: i64 = 0 87 while i < n1 { 88 let v: i64 = ta_constf((i % 5) - 2, 1) 89 ar[i] = v; ai[i] = TA_F32_ZERO; br[i] = v; bi[i] = TA_F32_ZERO 90 i = i + 1 91 } 92 let twr: *i64 = (sys_mmap(n1 * 8)) as *i64 93 let twi: *i64 = (sys_mmap(n1 * 8)) as *i64 94 fftf_twiddles(n1, twr, twi) 95 fftf_fwd(ar, ai, n1, twr, twi) 96 let dor: *i64 = (sys_mmap(n1 * 8)) as *i64 97 let doi: *i64 = (sys_mmap(n1 * 8)) as *i64 98 ft_dft(br, bi, n1, dor, doi) 99 var g1max: i64 = TA_F32_ZERO 100 i = 0 101 while i < n1 { 102 let d1: i64 = nx_f32_abs(nx_f32_sub(ar[i], dor[i])) 103 let d2: i64 = nx_f32_abs(nx_f32_sub(ai[i], doi[i])) 104 if nx_f32_gt(d1, g1max) == 1 { g1max = d1 } 105 if nx_f32_gt(d2, g1max) == 1 { g1max = d2 } 106 i = i + 1 107 } 108 let g1milli: i64 = ta_f32_to_milli(g1max) 109 if g1milli > 50 { ok = 0 } // FFT vs DFT agree to < 0.05 (f32 rounding-order only) 110 111 // ---------- G2: FNet-op gradcheck through the autograd ---------- 112 let x: *i64 = (sys_mmap(16 * 8)) as *i64 113 let target: *i64 = (sys_mmap(16 * 8)) as *i64 114 i = 0 115 while i < 16 { 116 x[i] = ta_constf((i % 5) - 2, 2) // -1 .. 1 117 target[i] = ta_constf((i % 7) - 3, 4) // -0.75 .. 0.75 118 i = i + 1 119 } 120 let ana: *i64 = (sys_mmap(16 * 8)) as *i64 121 g_fnet_grads(tape, vals, grads, st, x, target, ana) 122 let h: i64 = ta_constf(1, 128) 123 let flo: i64 = ta_constf(1, 64) 124 let tol: i64 = ta_constf(1, 32) 125 let xp: *i64 = (sys_mmap(16 * 8)) as *i64 126 let xm: *i64 = (sys_mmap(16 * 8)) as *i64 127 var g2pass: i64 = 1 128 var g2worst: i64 = 0 129 var c: i64 = 0 130 while c < 16 { 131 var j: i64 = 0 132 while j < 16 { xp[j] = x[j]; xm[j] = x[j]; j = j + 1 } 133 xp[c] = nx_f32_add(x[c], h) 134 xm[c] = nx_f32_sub(x[c], h) 135 let lp: i64 = g_fnet_loss(tape, vals, st, xp, target) 136 let lm: i64 = g_fnet_loss(tape, vals, st, xm, target) 137 let fd: i64 = nx_f32_div(nx_f32_sub(lp, lm), nx_f32_add(h, h)) 138 let num: i64 = nx_f32_abs(nx_f32_sub(fd, ana[c])) 139 var den: i64 = nx_f32_abs(ana[c]) 140 if nx_f32_lt(den, flo) == 1 { den = flo } 141 if nx_f32_lt(num, nx_f32_mul(tol, den)) != 1 { g2pass = 0 } 142 let nm: i64 = ta_f32_to_milli(num) 143 if nm > g2worst { g2worst = nm } 144 c = c + 1 145 } 146 if g2pass != 1 { ok = 0 } 147 148 // ---------- G3: train the input through the FNet op (loss must strictly decrease) ---------- 149 let xt: *i64 = (sys_mmap(16 * 8)) as *i64 150 i = 0 151 while i < 16 { xt[i] = TA_F32_ZERO; i = i + 1 } 152 let lr: i64 = ta_constf(1, 64) 153 let dxt: *i64 = (sys_mmap(16 * 8)) as *i64 154 var lfirst: i64 = TA_F32_ZERO 155 var llast: i64 = TA_F32_ZERO 156 var ep: i64 = 0 157 while ep < 1000 { 158 if ep == 0 { lfirst = g_fnet_loss(tape, vals, st, xt, target) } 159 g_fnet_grads(tape, vals, grads, st, xt, target, dxt) 160 i = 0 161 while i < 16 { xt[i] = nx_f32_sub(xt[i], nx_f32_mul(lr, dxt[i])); i = i + 1 } 162 ep = ep + 1 163 } 164 llast = g_fnet_loss(tape, vals, st, xt, target) 165 var g3pass: i64 = 1 166 if nx_f32_lt(llast, lfirst) != 1 { g3pass = 0 } // loss strictly decreased 167 if g3pass != 1 { ok = 0 } 168 169 // ---------- emit ---------- 170 var fdi: i64 = 1 171 while fdi >= 0 { 172 var out: i64 = 1 173 if fdi == 0 { out = sys_openat_append(FT_LOG, 420) } 174 if out >= 0 { 175 ft_w(out, "FNETTRAINGATE authored=organ block=trainable-fnet-f32" as *u8) 176 ft_w(out, " | G1_fft_vs_dft_milli=" as *u8); ft_wn(out, g1milli) 177 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) 178 ft_w(out, " | G3_train_pass=" as *u8); ft_wn(out, g3pass) 179 ft_w(out, " loss_first_milli=" as *u8); ft_wn(out, ta_f32_to_milli(lfirst)) 180 ft_w(out, " loss_last_milli=" as *u8); ft_wn(out, ta_f32_to_milli(llast)) 181 if ok == 1 { ft_w(out, " verdict=GREEN\n" as *u8) } else { ft_w(out, " verdict=RED\n" as *u8) } 182 if fdi == 0 { sys_close(out) } 183 } 184 fdi = fdi - 1 185 } 186 187 if ok == 1 { return 0 } 188 return 1 189}