`nx_vec_vq.nx` source

1// nx_vec_vq.nx -- R-VEC-4 of the onsite-search S-class ladder: SOVEREIGN vector quantization (LIBRARY). Raw f32
2// embedding vectors blow the RAM budget at 366k scale (Rule 21); VQ compresses each vector to ONE codebook index
3// -- "a codebook" of representative centroids (cited srch_vq.raw, Vector_quantization). k-means (Lloyd): assign
4// each vector to its nearest centroid, move each centroid to the mean of its members, repeat. Integer, no-float,
5// deterministic init (evenly-spaced seeds). HONEST SCOPE: single-codebook VQ; PRODUCT quantization (split the
6// vector into sub-spaces, a codebook per sub-space) is the named extension for finer compression.
7//
8// exports: vr_vq_train, vr_vq_encode, vr_vq_decode.   license_tier: ORIGINAL
9import "nx_syscalls.nx"
10
11// squared L2 distance
12func vq_l2(a: *i64, b: *i64, D: i64) -> i64 { var s: i64=0; var i: i64=0; while i<D { let d: i64=a[i]-b[i]; s=s+d*d; i=i+1 } return s }
13
14// encode: index of the nearest centroid (the compressed code for `vec`)
15func vr_vq_encode(vec: *i64, cb: *i64, K: i64, D: i64) -> i64 {
16    var best: i64=0; var bd: i64=vq_l2(vec, cb, D); var k: i64=1
17    while k<K { let ck: *i64=((cb as i64)+k*D*8) as *i64; let d: i64=vq_l2(vec, ck, D); if d<bd { bd=d; best=k } k=k+1 }
18    return best
19}
20
21// decode: reconstruct the vector as its centroid (pointer into the codebook)
22func vr_vq_decode(cb: *i64, idx: i64, D: i64) -> *i64 { return ((cb as i64)+idx*D*8) as *i64 }
23
24// train K centroids over N D-dim vectors with `iters` Lloyd iterations; writes the codebook cb[K*D].
25func vr_vq_train(V: *i64, N: i64, D: i64, K: i64, iters: i64, cb: *i64) -> i64 {
26    var k: i64=0
27    while k<K { let src: *i64=((V as i64)+(k*N/K)*D*8) as *i64; var d: i64=0; while d<D { cb[k*D+d]=src[d]; d=d+1 } k=k+1 }
28    let assign: *i64=sys_mmap(8*(N+2)) as *i64
29    let sum: *i64=sys_mmap(8*K*D) as *i64
30    let cnt: *i64=sys_mmap(8*(K+2)) as *i64
31    var it: i64=0
32    while it<iters {
33        var i: i64=0
34        while i<N { let vi: *i64=((V as i64)+i*D*8) as *i64; assign[i]=vr_vq_encode(vi, cb, K, D); i=i+1 }
35        var z: i64=0; while z<K*D { sum[z]=0; z=z+1 } z=0; while z<K { cnt[z]=0; z=z+1 }
36        i=0
37        while i<N { let a: i64=assign[i]; let vi: *i64=((V as i64)+i*D*8) as *i64; var d: i64=0; while d<D { sum[a*D+d]=sum[a*D+d]+vi[d]; d=d+1 } cnt[a]=cnt[a]+1; i=i+1 }
38        k=0; while k<K { if cnt[k]>0 { var d: i64=0; while d<D { cb[k*D+d]=sum[k*D+d]/cnt[k]; d=d+1 } } k=k+1 }
39        it=it+1
40    }
41    return 0
42}
nx_vec_vq.nx source

`nx_vec_vq.nx` source
