_tensor_grad_authored.nx source

1// AUTHORED BY THE NISHI BUILDER (nx_module_author tensor-grad template) -- T6 tensor autograd. 2// Tutoring: specs/2026-06-09-tutoring-training-substrate-rung2-tensor.md 3import "nx_syscalls.nx" 4import "nx_f32.nx" 5import "nx_f32_div.nx" 6import "nx_f32_cvt.nx" 7func _tg_puts(s: *u8) -> i64 { var n: i64=0; while s[n]!=(0 as u8){n=n+1} sys_write(1,s,n); return 0 } 8func _tg_num(v: i64) -> i64 { let bb: *u8=sys_mmap(28); var m: i64=v; if m<0{m=0-m;sys_write(1,"-" as *u8,1)}; let t: *u8=sys_mmap(28); var k: i64=0; if m==0{t[0]=48;k=1}; while m>0{t[k]=48+(m%10);m=m/10;k=k+1}; var i: i64=0; while i<k{bb[i]=t[k-1-i];i=i+1}; sys_write(1,bb,k); return 0 } 9func tg_q(n: i64, d: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(n), nx_i32_to_f32(d)) } 10func tg_milli(v: i64) -> i64 { 11 var neg: i64 = 0 12 var a: i64 = v 13 if nx_f32_lt(a, 0) == 1 { neg = 1; a = nx_f32_abs(a) } 14 let m: i64 = nx_f32_mul(a, nx_i32_to_f32(1000)) 15 var lo: i64 = 0 16 var hi: i64 = 16777216 17 while lo < hi { 18 let mid: i64 = (lo + hi + 1) / 2 19 if nx_f32_lt(m, nx_i32_to_f32(mid)) == 1 { hi = mid - 1 } else { lo = mid } 20 } 21 if neg == 1 { return 0 - lo } 22 return lo 23} 24const TG_NCAP: i64 = 256 25const TG_ACAP: i64 = 16384 26func tg_alloc(arena: *i64, ab: *i64, n: i64) -> *i64 { 27 if ab[0] + n > TG_ACAP { _tg_puts("ARENA OVERFLOW (honest fail)\n" as *u8); sys_exit(1) } 28 let base: i64 = arena as i64 29 let q: *i64 = (base + ab[0] * 8) as *i64 30 var i: i64 = 0 31 while i < n { q[i] = 0; i = i + 1 } 32 ab[0] = ab[0] + n 33 return q 34} 35func tg_node(tape: *i64, nb: *i64, op: i64, ai: i64, bi: i64, rows: i64, cols: i64, valp: *i64, gradp: *i64) -> i64 { 36 if nb[0] >= TG_NCAP { _tg_puts("TAPE OVERFLOW (honest fail)\n" as *u8); sys_exit(1) } 37 let r: i64 = nb[0] 38 tape[r*8+0] = op 39 tape[r*8+1] = ai 40 tape[r*8+2] = bi 41 tape[r*8+3] = rows 42 tape[r*8+4] = cols 43 tape[r*8+5] = valp as i64 44 tape[r*8+6] = gradp as i64 45 tape[r*8+7] = 0 46 nb[0] = r + 1 47 return r 48} 49func tg_valp(tape: *i64, x: i64) -> *i64 { return tape[x*8+5] as *i64 } 50func tg_gradp(tape: *i64, x: i64) -> *i64 { return tape[x*8+6] as *i64 } 51func tg_leaf(tape: *i64, nb: *i64, arena: *i64, ab: *i64, p: *i64, rows: i64, cols: i64) -> i64 { 52 return tg_node(tape, nb, 0, 0, 0, rows, cols, p, tg_alloc(arena, ab, rows*cols)) 53} 54func tg_matvec(tape: *i64, nb: *i64, arena: *i64, ab: *i64, a: i64, x: i64) -> i64 { 55 let r: i64 = tape[a*8+3] 56 let c: i64 = tape[a*8+4] 57 let av: *i64 = tg_valp(tape, a) 58 let xv: *i64 = tg_valp(tape, x) 59 let y: *i64 = tg_alloc(arena, ab, r) 60 var i: i64 = 0 61 while i < r { 62 var s: i64 = 0 63 var j: i64 = 0 64 while j < c { s = nx_f32_add(s, nx_f32_mul(av[i*c+j], xv[j])); j = j + 1 } 65 y[i] = s 66 i = i + 1 67 } 68 return tg_node(tape, nb, 1, a, x, r, 1, y, tg_alloc(arena, ab, r)) 69} 70func tg_addvec(tape: *i64, nb: *i64, arena: *i64, ab: *i64, x: i64, y: i64) -> i64 { 71 let n: i64 = tape[x*8+3] 72 let xv: *i64 = tg_valp(tape, x) 73 let yv: *i64 = tg_valp(tape, y) 74 let o: *i64 = tg_alloc(arena, ab, n) 75 var i: i64 = 0 76 while i < n { o[i] = nx_f32_add(xv[i], yv[i]); i = i + 1 } 77 return tg_node(tape, nb, 2, x, y, n, 1, o, tg_alloc(arena, ab, n)) 78} 79func tg_reluvec(tape: *i64, nb: *i64, arena: *i64, ab: *i64, x: i64) -> i64 { 80 let n: i64 = tape[x*8+3] 81 let xv: *i64 = tg_valp(tape, x) 82 let o: *i64 = tg_alloc(arena, ab, n) 83 var i: i64 = 0 84 while i < n { 85 var v: i64 = xv[i] 86 if nx_f32_lt(v, 0) == 1 { v = 0 } 87 o[i] = v 88 i = i + 1 89 } 90 return tg_node(tape, nb, 3, x, 0, n, 1, o, tg_alloc(arena, ab, n)) 91} 92func tg_mse(tape: *i64, nb: *i64, arena: *i64, ab: *i64, p: i64, t: i64) -> i64 { 93 let n: i64 = tape[p*8+3] 94 let pv: *i64 = tg_valp(tape, p) 95 let tv: *i64 = tg_valp(tape, t) 96 var acc: i64 = 0 97 var i: i64 = 0 98 while i < n { 99 let d: i64 = nx_f32_sub(pv[i], tv[i]) 100 acc = nx_f32_add(acc, nx_f32_mul(d, d)) 101 i = i + 1 102 } 103 let o: *i64 = tg_alloc(arena, ab, 1) 104 o[0] = nx_f32_mul(acc, tg_q(1, n)) 105 return tg_node(tape, nb, 4, p, t, 1, 1, o, tg_alloc(arena, ab, 1)) 106} 107func tg_smul(tape: *i64, nb: *i64, arena: *i64, ab: *i64, x: i64, y: i64) -> i64 { 108 let xv: *i64 = tg_valp(tape, x) 109 let yv: *i64 = tg_valp(tape, y) 110 let o: *i64 = tg_alloc(arena, ab, 1) 111 o[0] = nx_f32_mul(xv[0], yv[0]) 112 return tg_node(tape, nb, 5, x, y, 1, 1, o, tg_alloc(arena, ab, 1)) 113} 114func tg_backward(tape: *i64, nt: i64, root: i64) -> i64 { 115 var i: i64 = 0 116 while i < nt { 117 let gp: *i64 = tg_gradp(tape, i) 118 let n: i64 = tape[i*8+3] * tape[i*8+4] 119 var k: i64 = 0 120 while k < n { gp[k] = 0; k = k + 1 } 121 i = i + 1 122 } 123 let rg: *i64 = tg_gradp(tape, root) 124 rg[0] = nx_i32_to_f32(1) 125 var r: i64 = nt - 1 126 while r >= 0 { 127 let op: i64 = tape[r*8+0] 128 let ai: i64 = tape[r*8+1] 129 let bi: i64 = tape[r*8+2] 130 let g: *i64 = tg_gradp(tape, r) 131 if op == 1 { 132 let rr: i64 = tape[ai*8+3] 133 let cc: i64 = tape[ai*8+4] 134 let av: *i64 = tg_valp(tape, ai) 135 let ag: *i64 = tg_gradp(tape, ai) 136 let xv: *i64 = tg_valp(tape, bi) 137 let xg: *i64 = tg_gradp(tape, bi) 138 var i2: i64 = 0 139 while i2 < rr { 140 var j: i64 = 0 141 while j < cc { 142 ag[i2*cc+j] = nx_f32_add(ag[i2*cc+j], nx_f32_mul(g[i2], xv[j])) 143 xg[j] = nx_f32_add(xg[j], nx_f32_mul(av[i2*cc+j], g[i2])) 144 j = j + 1 145 } 146 i2 = i2 + 1 147 } 148 } 149 if op == 2 { 150 let n2: i64 = tape[r*8+3] 151 let ag2: *i64 = tg_gradp(tape, ai) 152 let bg2: *i64 = tg_gradp(tape, bi) 153 var k2: i64 = 0 154 while k2 < n2 { ag2[k2] = nx_f32_add(ag2[k2], g[k2]); bg2[k2] = nx_f32_add(bg2[k2], g[k2]); k2 = k2 + 1 } 155 } 156 if op == 3 { 157 let n3: i64 = tape[r*8+3] 158 let xv3: *i64 = tg_valp(tape, ai) 159 let ag3: *i64 = tg_gradp(tape, ai) 160 var k3: i64 = 0 161 while k3 < n3 { 162 if nx_f32_gt(xv3[k3], 0) == 1 { ag3[k3] = nx_f32_add(ag3[k3], g[k3]) } 163 k3 = k3 + 1 164 } 165 } 166 if op == 4 { 167 let n4: i64 = tape[ai*8+3] 168 let pv4: *i64 = tg_valp(tape, ai) 169 let tv4: *i64 = tg_valp(tape, bi) 170 let pg4: *i64 = tg_gradp(tape, ai) 171 let c2n: i64 = tg_q(2, n4) 172 var k4: i64 = 0 173 while k4 < n4 { 174 pg4[k4] = nx_f32_add(pg4[k4], nx_f32_mul(g[0], nx_f32_mul(c2n, nx_f32_sub(pv4[k4], tv4[k4])))) 175 k4 = k4 + 1 176 } 177 } 178 if op == 5 { 179 let av5: *i64 = tg_valp(tape, ai) 180 let bv5: *i64 = tg_valp(tape, bi) 181 let ag5: *i64 = tg_gradp(tape, ai) 182 let bg5: *i64 = tg_gradp(tape, bi) 183 ag5[0] = nx_f32_add(ag5[0], nx_f32_mul(g[0], bv5[0])) 184 bg5[0] = nx_f32_add(bg5[0], nx_f32_mul(g[0], av5[0])) 185 } 186 r = r - 1 187 } 188 return 0 189} 190func tg_fwd_mlp(w1: *i64, b1: *i64, w2: *i64, b2: *i64, x: *i64, t: i64) -> i64 { 191 let h: *i64 = sys_mmap(64) as *i64 192 var i: i64 = 0 193 while i < 2 { 194 var s: i64 = b1[i] 195 var j: i64 = 0 196 while j < 2 { s = nx_f32_add(s, nx_f32_mul(w1[i*2+j], x[j])); j = j + 1 } 197 if nx_f32_lt(s, 0) == 1 { s = 0 } 198 h[i] = s 199 i = i + 1 200 } 201 var o: i64 = b2[0] 202 i = 0 203 while i < 2 { o = nx_f32_add(o, nx_f32_mul(w2[i], h[i])); i = i + 1 } 204 let d: i64 = nx_f32_sub(o, t) 205 return nx_f32_mul(d, d) 206} 207func tg_pget(w1: *i64, b1: *i64, w2: *i64, b2: *i64, k: i64) -> i64 { 208 if k < 4 { return w1[k] } 209 if k < 6 { return b1[k-4] } 210 if k < 8 { return w2[k-6] } 211 return b2[0] 212} 213func tg_pset(w1: *i64, b1: *i64, w2: *i64, b2: *i64, k: i64, v: i64) -> i64 { 214 if k < 4 { w1[k] = v; return 0 } 215 if k < 6 { b1[k-4] = v; return 0 } 216 if k < 8 { w2[k-6] = v; return 0 } 217 b2[0] = v 218 return 0 219} 220func gate_a(tape: *i64, nb: *i64, arena: *i64, ab: *i64) -> i64 { 221 let w1: *i64 = sys_mmap(64) as *i64 222 let b1: *i64 = sys_mmap(64) as *i64 223 let w2: *i64 = sys_mmap(64) as *i64 224 let b2: *i64 = sys_mmap(64) as *i64 225 let x: *i64 = sys_mmap(64) as *i64 226 let t: *i64 = sys_mmap(64) as *i64 227 w1[0] = tg_q(2,3) 228 w1[1] = nx_f32_neg(tg_q(1,2)) 229 w1[2] = tg_q(1,4) 230 w1[3] = tg_q(1,2) 231 b1[0] = tg_q(1,4) 232 b1[1] = tg_q(1,8) 233 w2[0] = tg_q(3,2) 234 w2[1] = nx_f32_neg(tg_q(3,4)) 235 b2[0] = tg_q(1,8) 236 x[0] = tg_q(3,4) 237 x[1] = tg_q(1,2) 238 t[0] = tg_q(1,2) 239 nb[0] = 0 240 ab[0] = 0 241 let lw1: i64 = tg_leaf(tape, nb, arena, ab, w1, 2, 2) 242 let lb1: i64 = tg_leaf(tape, nb, arena, ab, b1, 2, 1) 243 let lw2: i64 = tg_leaf(tape, nb, arena, ab, w2, 1, 2) 244 let lb2: i64 = tg_leaf(tape, nb, arena, ab, b2, 1, 1) 245 let lx: i64 = tg_leaf(tape, nb, arena, ab, x, 2, 1) 246 let lt: i64 = tg_leaf(tape, nb, arena, ab, t, 1, 1) 247 let hh: i64 = tg_reluvec(tape, nb, arena, ab, tg_addvec(tape, nb, arena, ab, tg_matvec(tape, nb, arena, ab, lw1, lx), lb1)) 248 let oo: i64 = tg_addvec(tape, nb, arena, ab, tg_matvec(tape, nb, arena, ab, lw2, hh), lb2) 249 let loss: i64 = tg_mse(tape, nb, arena, ab, oo, lt) 250 tg_backward(tape, nb[0], loss) 251 let h128: i64 = tg_q(1,128) 252 var pass: i64 = 1 253 var k: i64 = 0 254 while k < 9 { 255 var ga: i64 = 0 256 if k < 4 { let gp: *i64 = tg_gradp(tape, lw1); ga = gp[k] } 257 else { if k < 6 { let gp: *i64 = tg_gradp(tape, lb1); ga = gp[k-4] } 258 else { if k < 8 { let gp: *i64 = tg_gradp(tape, lw2); ga = gp[k-6] } 259 else { let gp: *i64 = tg_gradp(tape, lb2); ga = gp[0] } } } 260 let save: i64 = tg_pget(w1, b1, w2, b2, k) 261 tg_pset(w1, b1, w2, b2, k, nx_f32_add(save, h128)) 262 let fp: i64 = tg_fwd_mlp(w1, b1, w2, b2, x, t[0]) 263 tg_pset(w1, b1, w2, b2, k, nx_f32_sub(save, h128)) 264 let fm: i64 = tg_fwd_mlp(w1, b1, w2, b2, x, t[0]) 265 tg_pset(w1, b1, w2, b2, k, save) 266 let fdif: i64 = nx_f32_div(nx_f32_sub(fp, fm), tg_q(1,64)) 267 var den: i64 = nx_f32_abs(fdif) 268 if nx_f32_lt(den, tg_q(1,64)) == 1 { den = tg_q(1,64) } 269 let rel: i64 = nx_f32_div(nx_f32_abs(nx_f32_sub(ga, fdif)), den) 270 _tg_puts(" gradcheck p" as *u8); _tg_num(k) 271 _tg_puts(" analytic-milli=" as *u8); _tg_num(tg_milli(ga)) 272 _tg_puts(" fd-milli=" as *u8); _tg_num(tg_milli(fdif)) 273 _tg_puts(" rel-milli=" as *u8); _tg_num(tg_milli(rel)); _tg_puts("\n" as *u8) 274 if nx_f32_lt(rel, tg_q(1,32)) == 0 { pass = 0 } 275 k = k + 1 276 } 277 return pass 278} 279func tg_train(tape: *i64, nb: *i64, arena: *i64, ab: *i64, out: *i64) -> i64 { 280 let at: *i64 = sys_mmap(64) as *i64 281 let ct: *i64 = sys_mmap(64) as *i64 282 at[0] = tg_q(3,2) 283 at[1] = nx_f32_neg(tg_q(1,2)) 284 at[2] = tg_q(1,4) 285 at[3] = nx_i32_to_f32(1) 286 ct[0] = nx_f32_neg(tg_q(1,2)) 287 ct[1] = tg_q(3,4) 288 let xs: *i64 = sys_mmap(256) as *i64 289 let ys: *i64 = sys_mmap(256) as *i64 290 var k: i64 = 0 291 while k < 8 { 292 xs[k*2+0] = tg_q(k, 4) 293 xs[k*2+1] = tg_q((k*3) % 8, 8) 294 var i: i64 = 0 295 while i < 2 { 296 ys[k*2+i] = nx_f32_add(nx_f32_add(nx_f32_mul(at[i*2+0], xs[k*2+0]), nx_f32_mul(at[i*2+1], xs[k*2+1])), ct[i]) 297 i = i + 1 298 } 299 k = k + 1 300 } 301 let lr: i64 = tg_q(1,10) 302 let inv8: *i64 = sys_mmap(64) as *i64 303 inv8[0] = tg_q(1,8) 304 let ww: *i64 = sys_mmap(64) as *i64 305 let bb2: *i64 = sys_mmap(64) as *i64 306 var first: i64 = 0 307 var loss_v: i64 = 0 308 var ep: i64 = 0 309 while ep < 400 { 310 nb[0] = 0 311 ab[0] = 0 312 let lw: i64 = tg_leaf(tape, nb, arena, ab, ww, 2, 2) 313 let lb: i64 = tg_leaf(tape, nb, arena, ab, bb2, 2, 1) 314 var accn: i64 = 0 - 1 315 k = 0 316 while k < 8 { 317 let xb: i64 = xs as i64 318 let yb: i64 = ys as i64 319 let lx: i64 = tg_leaf(tape, nb, arena, ab, (xb + k*16) as *i64, 2, 1) 320 let ly: i64 = tg_leaf(tape, nb, arena, ab, (yb + k*16) as *i64, 2, 1) 321 let pred: i64 = tg_addvec(tape, nb, arena, ab, tg_matvec(tape, nb, arena, ab, lw, lx), lb) 322 let m: i64 = tg_mse(tape, nb, arena, ab, pred, ly) 323 if accn < 0 { accn = m } else { accn = tg_addvec(tape, nb, arena, ab, accn, m) } 324 k = k + 1 325 } 326 let li: i64 = tg_leaf(tape, nb, arena, ab, inv8, 1, 1) 327 let loss: i64 = tg_smul(tape, nb, arena, ab, accn, li) 328 tg_backward(tape, nb[0], loss) 329 let lv: *i64 = tg_valp(tape, loss) 330 loss_v = lv[0] 331 if ep == 0 { first = loss_v } 332 let gw: *i64 = tg_gradp(tape, lw) 333 let gb: *i64 = tg_gradp(tape, lb) 334 var q2: i64 = 0 335 while q2 < 4 { ww[q2] = nx_f32_sub(ww[q2], nx_f32_mul(lr, gw[q2])); q2 = q2 + 1 } 336 q2 = 0 337 while q2 < 2 { bb2[q2] = nx_f32_sub(bb2[q2], nx_f32_mul(lr, gb[q2])); q2 = q2 + 1 } 338 if ep % 100 == 0 { _tg_puts(" epoch=" as *u8); _tg_num(ep); _tg_puts(" loss-milli=" as *u8); _tg_num(tg_milli(loss_v)); _tg_puts("\n" as *u8) } 339 ep = ep + 1 340 } 341 var i3: i64 = 0 342 while i3 < 4 { out[i3] = ww[i3]; i3 = i3 + 1 } 343 out[4] = bb2[0] 344 out[5] = bb2[1] 345 out[6] = loss_v 346 out[7] = first 347 return 0 348} 349func ad_step(w: *i64, g: *i64, m: *i64, v: *i64, n: i64, lr: i64, b1: i64, b2: i64, eps: i64, wd: i64, t: i64) -> i64 { 350 var c1: i64 = nx_i32_to_f32(1) 351 var c2: i64 = nx_i32_to_f32(1) 352 var k: i64 = 0 353 while k < t { c1 = nx_f32_mul(c1, b1); c2 = nx_f32_mul(c2, b2); k = k + 1 } 354 let bc1: i64 = nx_f32_sub(nx_i32_to_f32(1), c1) 355 let bc2: i64 = nx_f32_sub(nx_i32_to_f32(1), c2) 356 let omb1: i64 = nx_f32_sub(nx_i32_to_f32(1), b1) 357 let omb2: i64 = nx_f32_sub(nx_i32_to_f32(1), b2) 358 var i: i64 = 0 359 while i < n { 360 m[i] = nx_f32_add(nx_f32_mul(b1, m[i]), nx_f32_mul(omb1, g[i])) 361 v[i] = nx_f32_add(nx_f32_mul(b2, v[i]), nx_f32_mul(omb2, nx_f32_mul(g[i], g[i]))) 362 let mh: i64 = nx_f32_div(m[i], bc1) 363 let vh: i64 = nx_f32_div(v[i], bc2) 364 let upd: i64 = nx_f32_div(mh, nx_f32_add(nx_f32_sqrt(vh), eps)) 365 w[i] = nx_f32_sub(w[i], nx_f32_add(nx_f32_mul(lr, upd), nx_f32_mul(nx_f32_mul(lr, wd), w[i]))) 366 i = i + 1 367 } 368 return 0 369} 370func ad_train(tape: *i64, nb: *i64, arena: *i64, ab: *i64, out: *i64) -> i64 { 371 let at: *i64 = sys_mmap(64) as *i64 372 let ct: *i64 = sys_mmap(64) as *i64 373 at[0] = tg_q(3,2) 374 at[1] = nx_f32_neg(tg_q(1,2)) 375 at[2] = tg_q(1,4) 376 at[3] = nx_i32_to_f32(1) 377 ct[0] = nx_f32_neg(tg_q(1,2)) 378 ct[1] = tg_q(3,4) 379 let xs: *i64 = sys_mmap(256) as *i64 380 let ys: *i64 = sys_mmap(256) as *i64 381 var k: i64 = 0 382 while k < 8 { 383 xs[k*2+0] = tg_q(k, 4) 384 xs[k*2+1] = tg_q((k*3) % 8, 8) 385 var i: i64 = 0 386 while i < 2 { 387 ys[k*2+i] = nx_f32_add(nx_f32_add(nx_f32_mul(at[i*2+0], xs[k*2+0]), nx_f32_mul(at[i*2+1], xs[k*2+1])), ct[i]) 388 i = i + 1 389 } 390 k = k + 1 391 } 392 let lr: i64 = tg_q(1,20) 393 let b1c: i64 = tg_q(9,10) 394 let b2c: i64 = tg_q(999,1000) 395 let eps: i64 = tg_q(1,100000) 396 let inv8: *i64 = sys_mmap(64) as *i64 397 inv8[0] = tg_q(1,8) 398 let ww: *i64 = sys_mmap(64) as *i64 399 let bb2: *i64 = sys_mmap(64) as *i64 400 let mw: *i64 = sys_mmap(64) as *i64 401 let vw: *i64 = sys_mmap(64) as *i64 402 let mb: *i64 = sys_mmap(64) as *i64 403 let vb: *i64 = sys_mmap(64) as *i64 404 var first: i64 = 0 405 var loss_v: i64 = 0 406 var ep: i64 = 0 407 while ep < 400 { 408 nb[0] = 0 409 ab[0] = 0 410 let lw: i64 = tg_leaf(tape, nb, arena, ab, ww, 2, 2) 411 let lb: i64 = tg_leaf(tape, nb, arena, ab, bb2, 2, 1) 412 var accn: i64 = 0 - 1 413 k = 0 414 while k < 8 { 415 let xb: i64 = xs as i64 416 let yb: i64 = ys as i64 417 let lx: i64 = tg_leaf(tape, nb, arena, ab, (xb + k*16) as *i64, 2, 1) 418 let ly: i64 = tg_leaf(tape, nb, arena, ab, (yb + k*16) as *i64, 2, 1) 419 let pred: i64 = tg_addvec(tape, nb, arena, ab, tg_matvec(tape, nb, arena, ab, lw, lx), lb) 420 let m2: i64 = tg_mse(tape, nb, arena, ab, pred, ly) 421 if accn < 0 { accn = m2 } else { accn = tg_addvec(tape, nb, arena, ab, accn, m2) } 422 k = k + 1 423 } 424 let li: i64 = tg_leaf(tape, nb, arena, ab, inv8, 1, 1) 425 let loss: i64 = tg_smul(tape, nb, arena, ab, accn, li) 426 tg_backward(tape, nb[0], loss) 427 let lv: *i64 = tg_valp(tape, loss) 428 loss_v = lv[0] 429 if ep == 0 { first = loss_v } 430 ad_step(ww, tg_gradp(tape, lw), mw, vw, 4, lr, b1c, b2c, eps, 0, ep + 1) 431 ad_step(bb2, tg_gradp(tape, lb), mb, vb, 2, lr, b1c, b2c, eps, 0, ep + 1) 432 if ep % 100 == 0 { _tg_puts(" adamw epoch=" as *u8); _tg_num(ep); _tg_puts(" loss-milli=" as *u8); _tg_num(tg_milli(loss_v)); _tg_puts("\n" as *u8) } 433 ep = ep + 1 434 } 435 var i4: i64 = 0 436 while i4 < 4 { out[i4] = ww[i4]; i4 = i4 + 1 } 437 out[4] = bb2[0] 438 out[5] = bb2[1] 439 out[6] = loss_v 440 out[7] = first 441 return 0 442} 443func main() -> i64 { 444 _tg_puts("=== BUILDER-AUTHORED TENSOR AUTOGRAD: T6 (matvec/relu/mse identities, arena, 9-param gradcheck) ===\n" as *u8) 445 let tape: *i64 = sys_mmap(32768) as *i64 446 let nb: *i64 = sys_mmap(16) as *i64 447 let arena: *i64 = sys_mmap(131072) as *i64 448 let ab: *i64 = sys_mmap(16) as *i64 449 let pa: i64 = gate_a(tape, nb, arena, ab) 450 if pa == 1 { _tg_puts(" GATE A MLP gradcheck (9 params): PASS\n" as *u8) } else { _tg_puts(" GATE A MLP gradcheck: FAIL\n" as *u8) } 451 let o1: *i64 = sys_mmap(128) as *i64 452 let o2: *i64 = sys_mmap(128) as *i64 453 tg_train(tape, nb, arena, ab, o1) 454 _tg_puts(" run1 W-milli: " as *u8) 455 var k: i64 = 0 456 while k < 4 { _tg_num(tg_milli(o1[k])); _tg_puts(" " as *u8); k = k + 1 } 457 _tg_puts(" b-milli: " as *u8); _tg_num(tg_milli(o1[4])); _tg_puts(" " as *u8); _tg_num(tg_milli(o1[5])) 458 _tg_puts(" loss-milli=" as *u8); _tg_num(tg_milli(o1[6])); _tg_puts("\n" as *u8) 459 let at2: *i64 = sys_mmap(64) as *i64 460 at2[0] = tg_q(3,2) 461 at2[1] = nx_f32_neg(tg_q(1,2)) 462 at2[2] = tg_q(1,4) 463 at2[3] = nx_i32_to_f32(1) 464 at2[4] = nx_f32_neg(tg_q(1,2)) 465 at2[5] = tg_q(3,4) 466 var pb: i64 = 1 467 if nx_f32_lt(o1[6], tg_q(1,1000)) == 0 { pb = 0 } 468 if nx_f32_lt(o1[6], o1[7]) == 0 { pb = 0 } 469 k = 0 470 while k < 6 { 471 if nx_f32_lt(nx_f32_abs(nx_f32_sub(o1[k], at2[k])), tg_q(1,16)) == 0 { pb = 0 } 472 k = k + 1 473 } 474 if pb == 1 { _tg_puts(" GATE B affine map recovered (all 6 cells within 1/16, loss < 1/1000): PASS\n" as *u8) } else { _tg_puts(" GATE B: FAIL\n" as *u8) } 475 tg_train(tape, nb, arena, ab, o2) 476 var pc: i64 = 1 477 k = 0 478 while k < 7 { if o1[k] != o2[k] { pc = 0 } k = k + 1 } 479 if pc == 1 { _tg_puts(" GATE C bit-exact tensor training: PASS (the EXCEED axis holds at rung 2)\n" as *u8) } else { _tg_puts(" GATE C: FAIL\n" as *u8) } 480 let aw: *i64 = sys_mmap(64) as *i64 481 let ag2: *i64 = sys_mmap(64) as *i64 482 let am: *i64 = sys_mmap(64) as *i64 483 let av2: *i64 = sys_mmap(64) as *i64 484 aw[0] = nx_i32_to_f32(1) 485 aw[1] = nx_i32_to_f32(0-2) 486 aw[2] = nx_i32_to_f32(3) 487 aw[3] = nx_i32_to_f32(0-4) 488 ag2[0] = nx_i32_to_f32(2) 489 ag2[1] = nx_i32_to_f32(0-3) 490 ag2[2] = tg_q(1,2) 491 ag2[3] = nx_i32_to_f32(0-5) 492 let wprev: *i64 = sys_mmap(64) as *i64 493 k = 0 494 while k < 4 { wprev[k] = aw[k]; am[k] = 0; av2[k] = 0; k = k + 1 } 495 let alr: i64 = tg_q(1,100) 496 ad_step(aw, ag2, am, av2, 4, alr, tg_q(9,10), tg_q(999,1000), tg_q(1,100000), 0, 1) 497 var pd: i64 = 1 498 k = 0 499 while k < 4 { 500 var want: i64 = nx_f32_sub(wprev[k], alr) 501 if nx_f32_lt(ag2[k], 0) == 1 { want = nx_f32_add(wprev[k], alr) } 502 if nx_f32_lt(nx_f32_abs(nx_f32_sub(aw[k], want)), tg_q(1,2000)) == 0 { pd = 0 } 503 k = k + 1 504 } 505 if pd == 1 { _tg_puts(" GATE D AdamW first-step sign property (|step| == lr per param): PASS\n" as *u8) } else { _tg_puts(" GATE D: FAIL\n" as *u8) } 506 aw[0] = nx_i32_to_f32(1) 507 am[0] = 0 508 av2[0] = 0 509 ag2[0] = 0 510 let dlr: i64 = tg_q(1,10) 511 let dwd: i64 = tg_q(1,10) 512 var st: i64 = 0 513 while st < 10 { ad_step(aw, ag2, am, av2, 1, dlr, tg_q(9,10), tg_q(999,1000), tg_q(1,100000), dwd, st + 1); st = st + 1 } 514 var expw: i64 = nx_i32_to_f32(1) 515 let fac: i64 = nx_f32_sub(nx_i32_to_f32(1), nx_f32_mul(dlr, dwd)) 516 st = 0 517 while st < 10 { expw = nx_f32_mul(expw, fac); st = st + 1 } 518 var pe: i64 = 1 519 if nx_f32_lt(nx_f32_abs(nx_f32_sub(aw[0], expw)), tg_q(1,1000)) == 0 { pe = 0 } 520 if am[0] != 0 { pe = 0 } 521 if av2[0] != 0 { pe = 0 } 522 _tg_puts(" decay-only w-milli=" as *u8); _tg_num(tg_milli(aw[0])) 523 _tg_puts(" expected-milli=" as *u8); _tg_num(tg_milli(expw)); _tg_puts("\n" as *u8) 524 if pe == 1 { _tg_puts(" GATE E decoupled decay exact ((1-lr*wd)^k, moments untouched): PASS\n" as *u8) } else { _tg_puts(" GATE E: FAIL\n" as *u8) } 525 let o3: *i64 = sys_mmap(128) as *i64 526 let o4: *i64 = sys_mmap(128) as *i64 527 ad_train(tape, nb, arena, ab, o3) 528 _tg_puts(" adamw W-milli: " as *u8) 529 k = 0 530 while k < 4 { _tg_num(tg_milli(o3[k])); _tg_puts(" " as *u8); k = k + 1 } 531 _tg_puts(" b-milli: " as *u8); _tg_num(tg_milli(o3[4])); _tg_puts(" " as *u8); _tg_num(tg_milli(o3[5])) 532 _tg_puts(" loss-milli=" as *u8); _tg_num(tg_milli(o3[6])); _tg_puts("\n" as *u8) 533 var pf: i64 = 1 534 if nx_f32_lt(o3[6], tg_q(1,1000)) == 0 { pf = 0 } 535 if nx_f32_lt(o3[6], o3[7]) == 0 { pf = 0 } 536 k = 0 537 while k < 6 { 538 if nx_f32_lt(nx_f32_abs(nx_f32_sub(o3[k], at2[k])), tg_q(1,16)) == 0 { pf = 0 } 539 k = k + 1 540 } 541 ad_train(tape, nb, arena, ab, o4) 542 k = 0 543 while k < 7 { if o3[k] != o4[k] { pf = 0 } k = k + 1 } 544 if pf == 1 { _tg_puts(" GATE F AdamW trains (recovered + bit-exact): PASS\n" as *u8) } else { _tg_puts(" GATE F: FAIL\n" as *u8) } 545 if pa == 1 { if pb == 1 { if pc == 1 { if pd == 1 { if pe == 1 { if pf == 1 { 546 _tg_puts(" TENSOR-AUTOGRAD T6 + ADAMW T5 GATE: PASS (verified matrix gradients + the modern optimizer, bits-up)\n" as *u8) 547 sys_exit(0) 548 return 0 549 } } } } } } 550 _tg_puts(" TENSOR-AUTOGRAD/ADAMW GATE: FAIL\n" as *u8) 551 sys_exit(1) 552 return 1 553}