nx_vit_encoder_layer_gate.nx source
↩ module page · 80 lines · 5397 B
1// nx_vit_encoder_layer_gate.nx -- validate the ViT encoder block WITHOUT a PyTorch reference: (E1) attention core on
2// a hand-computable case (identity projections, 2 tokens, 1 head -> softmax([1/sqrt2,0]) mix), (E2) the EXACT
3// residual identity (all attn/MLP weights zero, LN gamma=1 beta=0 -> out == x, proving the residual wiring), (E3)
4// end-to-end finiteness with non-trivial weights. expect_exit: 0
5import "nx_syscalls.nx"
6import "nx_f32.nx"
7import "nx_f32_cvt.nx"
8import "nx_f32_div.nx"
9import "nx_vit_encoder_layer.nx"
10
11func gp(s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } return sys_write(1, s, n) }
12func gn(v: i64) -> i64 { var m: i64=v; if m<0{gp("-" as *u8);m=0-m} let t:*u8=sys_mmap(24); var k:i64=0; if m==0{t[0]=48 as u8;k=1} while m>0{t[k]=(48+(m%10)) as u8;m=m/10;k=k+1} var i:i64=0; let o:*u8=sys_mmap(24); while i<k{o[i]=t[k-1-i];i=i+1} sys_write(1,o,k); return 0 }
13func 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 }
14func zero(p: *i64, n: i64) -> i64 { var i: i64=0; while i<n { p[i]=0; i=i+1 } return 0 }
15func ones(p: *i64, n: i64) -> i64 { var i: i64=0; while i<n { p[i]=nx_i32_to_f32(1); i=i+1 } return 0 }
16
17func main(argc: i64, argv: *i64) -> i64 {
18 var pass: i64 = 0
19 let tol: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(100))
20 let eps: i64 = nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(1000000))
21
22 // ===== E1: attention core, identity projections, T=2 D=2 nh=1 hd=2 =====
23 let x2: *i64 = sys_mmap(8*4) as *i64
24 x2[0]=nx_i32_to_f32(1); x2[1]=0; x2[2]=0; x2[3]=nx_i32_to_f32(1) // [[1,0],[0,1]]
25 let I2: *i64 = sys_mmap(8*4) as *i64
26 I2[0]=nx_i32_to_f32(1); I2[1]=0; I2[2]=0; I2[3]=nx_i32_to_f32(1) // identity [out,in]
27 let z2: *i64 = sys_mmap(8*2) as *i64; zero(z2, 2)
28 let sc2: i64 = nx_f32_div(nx_i32_to_f32(1), nx_f32_sqrt(nx_i32_to_f32(2)))
29 let o2: *i64 = sys_mmap(8*4) as *i64
30 vit_mhsa(x2, I2,z2, I2,z2, I2,z2, I2,z2, 2, 2, 1, 2, sc2, o2)
31 let a: i64 = nx_f32_div(nx_i32_to_f32(6698), nx_i32_to_f32(10000))
32 let b: i64 = nx_f32_div(nx_i32_to_f32(3302), nx_i32_to_f32(10000))
33 if close(o2[0], a, tol) == 1 { if close(o2[1], b, tol) == 1 { if close(o2[2], b, tol) == 1 { if close(o2[3], a, tol) == 1 {
34 pass = pass + 1; gp("E1 attention core softmax mix [0.67,0.33] OK\n" as *u8)
35 } } } }
36 if close(o2[0], a, tol) == 0 { gp("E1 FAIL o2[0..3]\n" as *u8) }
37
38 // ===== E2/E3: full layer, T=4 D=8 nh=2 hd=4 mlp=16 =====
39 let T: i64=4; let D: i64=8; let MLP: i64=16
40 let x8: *i64 = sys_mmap(8*T*D) as *i64
41 var i: i64=0; while i<T*D { x8[i]=nx_i32_to_f32((i%5)-2); i=i+1 } // varied input
42 // tensors
43 let ln1g: *i64=sys_mmap(8*D) as *i64; let ln1b: *i64=sys_mmap(8*D) as *i64
44 let Wq: *i64=sys_mmap(8*D*D) as *i64; let bq: *i64=sys_mmap(8*D) as *i64
45 let Wk: *i64=sys_mmap(8*D*D) as *i64; let bk: *i64=sys_mmap(8*D) as *i64
46 let Wv: *i64=sys_mmap(8*D*D) as *i64; let bv: *i64=sys_mmap(8*D) as *i64
47 let Wo: *i64=sys_mmap(8*D*D) as *i64; let bo: *i64=sys_mmap(8*D) as *i64
48 let ln2g: *i64=sys_mmap(8*D) as *i64; let ln2b: *i64=sys_mmap(8*D) as *i64
49 let fc1w: *i64=sys_mmap(8*MLP*D) as *i64; let fc1b: *i64=sys_mmap(8*MLP) as *i64
50 let fc2w: *i64=sys_mmap(8*D*MLP) as *i64; let fc2b: *i64=sys_mmap(8*D) as *i64
51 let wts: *i64=sys_mmap(8*16) as *i64
52 wts[0]=ln1g as i64; wts[1]=ln1b as i64; wts[2]=Wq as i64; wts[3]=bq as i64; wts[4]=Wk as i64; wts[5]=bk as i64
53 wts[6]=Wv as i64; wts[7]=bv as i64; wts[8]=Wo as i64; wts[9]=bo as i64; wts[10]=ln2g as i64; wts[11]=ln2b as i64
54 wts[12]=fc1w as i64; wts[13]=fc1b as i64; wts[14]=fc2w as i64; wts[15]=fc2b as i64
55
56 // E2: zero everything, LN gammas=1 -> out == x exactly
57 ones(ln1g,D); zero(ln1b,D); ones(ln2g,D); zero(ln2b,D)
58 zero(Wq,D*D); zero(bq,D); zero(Wk,D*D); zero(bk,D); zero(Wv,D*D); zero(bv,D); zero(Wo,D*D); zero(bo,D)
59 zero(fc1w,MLP*D); zero(fc1b,MLP); zero(fc2w,D*MLP); zero(fc2b,D)
60 let sc4: i64 = nx_f32_div(nx_i32_to_f32(1), nx_f32_sqrt(nx_i32_to_f32(4)))
61 let oE: *i64 = sys_mmap(8*T*D) as *i64
62 vit_encoder_layer(x8, wts, T, D, 2, 4, MLP, eps, sc4, oE)
63 var ident: i64=1; i=0; while i<T*D { if oE[i] != x8[i] { ident=0 } i=i+1 }
64 if ident==1 { pass = pass + 1; gp("E2 zero-weights -> out == x EXACT (residual wiring) OK\n" as *u8) } else { gp("E2 FAIL not identity\n" as *u8) }
65
66 // E3: non-trivial weights -> finite + differs from x
67 i=0; while i<D*D { Wq[i]=nx_i32_to_f32((i%3)-1); Wk[i]=nx_i32_to_f32((i%2)); Wv[i]=nx_i32_to_f32((i%3)-1); Wo[i]=nx_i32_to_f32((i%2)); i=i+1 }
68 i=0; while i<MLP*D { fc1w[i]=nx_i32_to_f32((i%3)-1); i=i+1 }
69 i=0; while i<D*MLP { fc2w[i]=nx_i32_to_f32((i%2)); i=i+1 }
70 vit_encoder_layer(x8, wts, T, D, 2, 4, MLP, eps, sc4, oE)
71 var bad: i64=0; var diff: i64=0; i=0
72 while i<T*D { if nx_f32_is_nan(oE[i])==1{bad=bad+1} if nx_f32_is_inf(oE[i])==1{bad=bad+1} if oE[i]!=x8[i]{diff=diff+1} i=i+1 }
73 if bad==0 { if diff>0 { pass = pass + 1; gp("E3 non-trivial weights -> finite + transformed OK\n" as *u8) } }
74 if bad!=0 { gp("E3 FAIL bad=" as *u8); gn(bad); gp("\n" as *u8) }
75
76 gp("VIT-ENCODER-GATE pass=" as *u8); gn(pass); gp("/3\n" as *u8)
77 if pass == 3 { gp("VIT-ENCODER-GATE GREEN 3/3 (attention core + residual identity + end-to-end)\n" as *u8); sys_exit(0) }
78 sys_exit(1)
79 return 0
80}