nx_f32_bn_fold_gate.nx source
↩ module page · 60 lines · 3087 B
1// nx_f32_bn_fold_gate.nx -- proof of BN-fold: exact folded W'/b', AND functional equivalence foldedconv(in) ==
2// BN(conv(in)). Values chosen so var+eps=4 -> sqrt=2 (exact), s=gamma/2=3. expect_exit: 0
3import "nx_syscalls.nx"
4import "nx_f32_cvt.nx"
5import "nx_f32.nx"
6import "nx_f32_conv2d.nx"
7import "nx_f32_bn_fold.nx"
8
9func gp(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return sys_write(1, s, n) }
10func gn(v: i64) -> i64 {
11 let bb: *u8 = sys_mmap(28); var m: i64 = v
12 if m < 0 { sys_write(1, "-" as *u8, 1); m = 0 - m }
13 let t: *u8 = sys_mmap(28); var k: i64 = 0
14 if m == 0 { t[0] = 48 as u8; k = 1 }
15 while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
16 var i: i64 = 0; while i < k { bb[i] = t[k - 1 - i]; i = i + 1 }
17 return sys_write(1, bb, k)
18}
19
20func main(argc: i64, argv: *i64) -> i64 {
21 var pass: i64 = 0
22 // conv: C_out=1, C_in=2, 1x1 ; W=[1,2], b=[10] ; BN: gamma=6, beta=5, mean=1, var=3, eps=1 -> sqrt(4)=2, s=3
23 let W: *i64 = sys_mmap(8 * 4) as *i64
24 let b: *i64 = sys_mmap(8) as *i64
25 let gam: *i64 = sys_mmap(8) as *i64; let bet: *i64 = sys_mmap(8) as *i64
26 let men: *i64 = sys_mmap(8) as *i64; let vr: *i64 = sys_mmap(8) as *i64
27 let Wo: *i64 = sys_mmap(8 * 4) as *i64; let bo: *i64 = sys_mmap(8) as *i64
28 W[0]=nx_i32_to_f32(1); W[1]=nx_i32_to_f32(2); b[0]=nx_i32_to_f32(10)
29 gam[0]=nx_i32_to_f32(6); bet[0]=nx_i32_to_f32(5); men[0]=nx_i32_to_f32(1); vr[0]=nx_i32_to_f32(3)
30 let eps: i64 = nx_i32_to_f32(1)
31
32 nx_f32_bn_fold(W, b, 1, 2, gam, bet, men, vr, eps, Wo, bo)
33
34 // B1 exact folded weights W'=[3,6], bias b'=(10-1)*3+5=32
35 if Wo[0]==nx_i32_to_f32(3) { if Wo[1]==nx_i32_to_f32(6) { if bo[0]==nx_i32_to_f32(32) {
36 pass = pass + 1; gp("B1 fold W'=[3,6] b'=32 exact OK\n" as *u8)
37 } } }
38 if Wo[0]!=nx_i32_to_f32(3) { gp("B1 FAIL W0'=" as *u8); gn(Wo[0]); gp(" b'=" as *u8); gn(bo[0]); gp("\n" as *u8) }
39
40 // B2 functional equivalence: BN(conv(in)) == foldedconv(in) for in=[2,3]
41 let inp: *i64 = sys_mmap(8 * 2) as *i64
42 inp[0]=nx_i32_to_f32(2); inp[1]=nx_i32_to_f32(3)
43 let xr: *i64 = sys_mmap(8) as *i64; let yf: *i64 = sys_mmap(8) as *i64
44 // path A: conv then BN by hand -> x=1*2+2*3+10=18 ; BN=6*(18-1)/2+5 = 56
45 nx_f32_conv2d_forward(inp, 1, 2, 1, 1, W, 1, 1, 1, 1, 0, b, xr)
46 let denom: i64 = nx_f32_sqrt(nx_f32_add(vr[0], eps))
47 let s: i64 = nx_f32_div(gam[0], denom)
48 let bn: i64 = nx_f32_add(nx_f32_mul(nx_f32_sub(xr[0], men[0]), s), bet[0])
49 // path B: folded conv
50 nx_f32_conv2d_forward(inp, 1, 2, 1, 1, Wo, 1, 1, 1, 1, 0, bo, yf)
51 if bn == yf[0] { if yf[0]==nx_i32_to_f32(56) {
52 pass = pass + 1; gp("B2 foldedconv == BN(conv) == 56 OK\n" as *u8)
53 } }
54 if bn != yf[0] { gp("B2 FAIL bn=" as *u8); gn(bn); gp(" folded=" as *u8); gn(yf[0]); gp("\n" as *u8) }
55
56 gp("BN-FOLD-GATE pass=" as *u8); gn(pass); gp("/2\n" as *u8)
57 if pass == 2 { gp("BN-FOLD-GATE GREEN 2/2 (exact fold + functional equivalence conv+BN==foldedconv)\n" as *u8); sys_exit(0) }
58 sys_exit(1)
59 return 0
60}