nx_f32_layernorm_lib_gate.nx source
↩ module page · 55 lines · 3075 B
1// nx_f32_layernorm_lib_gate.nx -- proof of the flat LayerNorm vs hand-computed values + affine + 2-row. expect_exit: 0
2import "nx_syscalls.nx"
3import "nx_f32.nx"
4import "nx_f32_cvt.nx"
5import "nx_f32_div.nx"
6import "nx_f32_layernorm.nx"
7
8func gp(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return sys_write(1, s, n) }
9func close(a: i64, b: i64, tol: i64) -> i64 { if nx_f32_lt(nx_f32_abs(nx_f32_sub(a, b)), tol) == 1 { return 1 } return 0 }
10
11func main(argc: i64, argv: *i64) -> i64 {
12 var pass: i64 = 0
13 let tol: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(100)) // 0.01
14 let eps: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(1000000))
15 let x: *i64 = sys_mmap(8 * 8) as *i64
16 let g: *i64 = sys_mmap(8 * 4) as *i64
17 let b: *i64 = sys_mmap(8 * 4) as *i64
18 let o: *i64 = sys_mmap(8 * 8) as *i64
19
20 // row [1,2,3,4], gamma=1, beta=0 -> mean 2.5, std sqrt(1.25)=1.118034 -> [-1.3416,-0.4472,0.4472,1.3416]
21 x[0]=nx_i32_to_f32(1); x[1]=nx_i32_to_f32(2); x[2]=nx_i32_to_f32(3); x[3]=nx_i32_to_f32(4)
22 var i: i64=0; while i<4 { g[i]=nx_i32_to_f32(1); b[i]=nx_i32_to_f32(0); i=i+1 }
23 nx_f32_layernorm(x, g, b, 1, 4, eps, o)
24 let e0: i64 = nx_f32_neg(nx_f32_div(nx_i32_to_f32(13416), nx_i32_to_f32(10000)))
25 let e3: i64 = nx_f32_div(nx_i32_to_f32(13416), nx_i32_to_f32(10000))
26 if close(o[0], e0, tol) == 1 { if close(o[3], e3, tol) == 1 {
27 pass = pass + 1; gp("L1 LN([1,2,3,4]) -> [-1.34..,..,..,1.34] OK\n" as *u8)
28 } }
29 if close(o[0], e0, tol) == 0 { gp("L1 FAIL\n" as *u8) }
30
31 // property: normalized row has mean ~0
32 var sm: i64 = 0; i=0; while i<4 { sm = nx_f32_add(sm, o[i]); i=i+1 }
33 if close(sm, nx_i32_to_f32(0), tol) == 1 { pass = pass + 1; gp("L2 normalized mean ~0 OK\n" as *u8) } else { gp("L2 FAIL\n" as *u8) }
34
35 // affine: gamma=2, beta=1 -> out = xhat*2 + 1 ; o[3] ~ 1.3416*2+1 = 3.6833
36 i=0; while i<4 { g[i]=nx_i32_to_f32(2); b[i]=nx_i32_to_f32(1); i=i+1 }
37 nx_f32_layernorm(x, g, b, 1, 4, eps, o)
38 let ea: i64 = nx_f32_div(nx_i32_to_f32(36833), nx_i32_to_f32(10000))
39 if close(o[3], ea, tol) == 1 { pass = pass + 1; gp("L3 affine gamma=2 beta=1 OK\n" as *u8) } else { gp("L3 FAIL\n" as *u8) }
40
41 // 2-row independence: row0=[1,2,3,4], row1=[10,10,10,10] (var 0 -> xhat 0 -> out=beta=0 with g=1,b=0)
42 i=0; while i<4 { g[i]=nx_i32_to_f32(1); b[i]=nx_i32_to_f32(0); i=i+1 }
43 x[4]=nx_i32_to_f32(10); x[5]=nx_i32_to_f32(10); x[6]=nx_i32_to_f32(10); x[7]=nx_i32_to_f32(10)
44 nx_f32_layernorm(x, g, b, 2, 4, eps, o)
45 if close(o[0], e0, tol) == 1 { if close(o[4], nx_i32_to_f32(0), tol) == 1 {
46 pass = pass + 1; gp("L4 per-row independent (constant row -> ~0) OK\n" as *u8)
47 } }
48 if close(o[4], nx_i32_to_f32(0), tol) == 0 { gp("L4 FAIL\n" as *u8) }
49
50 gp("LAYERNORM-LIB-GATE pass=" as *u8)
51 let bb: *u8 = sys_mmap(8); bb[0]=(48+pass) as u8; sys_write(1,bb,1); gp("/4\n" as *u8)
52 if pass == 4 { gp("LAYERNORM-LIB-GATE GREEN 4/4 (exact + mean~0 + affine + per-row)\n" as *u8); sys_exit(0) }
53 sys_exit(1)
54 return 0
55}