nx_gen_archprobe.nx source
↩ module page · 53 lines · 2087 B
1// nx_gen_archprobe.nx -- print a DiT's architecture as READ FROM THE MODEL FILE.
2//
3// The hot-swap smoke test: point it at any checkpoint and see whether the engine can describe it
4// without being told anything. If this prints a coherent architecture, the same weights can drive
5// the sovereign DiT; if it refuses, the model needs work before it can be swapped in.
6//
7// Usage: nx_gen_archprobe <model_file>
8// license_tier: ORIGINAL
9
10import "nx_syscalls.nx"
11import "nx_le.nx"
12import "nx_f32.nx"
13import "nx_f32_div.nx"
14import "nx_f32_cvt.nx"
15import "nx_f16.nx"
16import "nx_strconv.nx"
17import "nx_genver.nx"
18import "nx_genweights.nx"
19import "nx_genarch.nx"
20
21func ap_puts(s: *u8) -> i64 {
22 var n: i64 = 0
23 while s[n] != (0 as u8) { n = n + 1 }
24 return sys_write(1, s, n)
25}
26
27func main(argc: i64, argv: *i64) -> i64 {
28 if argc < 2 {
29 ap_puts("usage: nx_gen_archprobe <model_file>\n" as *u8)
30 return 2
31 }
32 let gw: *i64 = nx_gw_open(argv[1] as *u8)
33 if (gw as i64) == 0 { ap_puts("open/parse failed (not a GGUF this build can read)\n" as *u8); return 20 }
34 nx_genver_emit("tensors" as *u8, nx_gw_ntensors(gw))
35
36 let arch: *i64 = sys_mmap(NX_ARCH_SLOTS * 8 + 64) as *i64
37 let rc: i64 = nx_arch_probe(gw, arch)
38 nx_genver_emit("probe_rc" as *u8, rc)
39 if rc != 0 { ap_puts("architecture not derivable from this file\n" as *u8); return 21 }
40
41 nx_genver_emit("hidden_dim" as *u8, arch[NX_ARCH_D])
42 nx_genver_emit("head_dim" as *u8, arch[NX_ARCH_HEAD_DIM])
43 nx_genver_emit("n_heads" as *u8, arch[NX_ARCH_N_HEADS])
44 nx_genver_emit("qkv_width" as *u8, arch[NX_ARCH_QKV_W])
45 nx_genver_emit("ffn_dim" as *u8, arch[NX_ARCH_FFN_DIM])
46 nx_genver_emit("n_layers" as *u8, arch[NX_ARCH_N_LAYERS])
47 nx_genver_emit("n_refiner_layers" as *u8, arch[NX_ARCH_N_REFINE])
48 nx_genver_emit("adaln_embed" as *u8, arch[NX_ARCH_ADALN_EMB])
49 nx_genver_emit("adaln_chunks" as *u8, arch[NX_ARCH_N_CHUNKS])
50 nx_genver_emit("out_dim" as *u8, arch[NX_ARCH_OUT_DIM])
51 nx_genver_emit("x_embed_in" as *u8, arch[NX_ARCH_XEMB_IN])
52 return 0
53}