nx_rvq.nx source
↩ module page · 46 lines · 2074 B
1// nx_rvq.nx -- RESIDUAL VECTOR QUANTIZATION, moonshot rung #3 and the literal structure inside SoundStream / Lyra-v2 /
2// Encodec: a stack of S codebooks where stage k quantises the RESIDUAL left by stages 0..k-1. Reconstruction is the sum
3// of the chosen stage centroids; each added stage costs log2(K) bits and shrinks the residual -> a rate-distortion knob.
4// The neural codec replaces these fixed centroids with learned ones over a learned latent; the structure is identical.
5// Composes nx_vq (+ nx_vq_train trains each stage on its residuals). license_tier: ORIGINAL
6import "nx_vq.nx"
7
8// pointer to stage k's codebook inside the flat cbs[S*K*D] buffer
9func rvq_stage(cbs: *i64, k: i64, K: i64, D: i64) -> *i64 { return (cbs as i64 + k*K*D*8) as *i64 }
10
11// encode v through S stages; writes indices[S]; resid (D) is caller scratch (holds the final residual on return).
12func rvq_encode(v: *i64, D: i64, cbs: *i64, K: i64, S: i64, indices: *i64, resid: *i64) -> i64 {
13 var i: i64 = 0
14 while i < D { resid[i] = v[i]; i = i + 1 }
15 var k: i64 = 0
16 while k < S {
17 let scb: *i64 = rvq_stage(cbs, k, K, D)
18 let idx: i64 = vq_encode(resid, scb, K, D)
19 indices[k] = idx
20 var j: i64 = 0
21 while j < D { resid[j] = resid[j] - scb[idx*D + j]; j = j + 1 }
22 k = k + 1
23 }
24 return 0
25}
26// reconstruct using the first `stages` indices -> out[D] = sum of those stage centroids
27func rvq_decode_partial(indices: *i64, cbs: *i64, K: i64, D: i64, stages: i64, out: *i64) -> i64 {
28 var i: i64 = 0
29 while i < D { out[i] = 0; i = i + 1 }
30 var k: i64 = 0
31 while k < stages {
32 let scb: *i64 = rvq_stage(cbs, k, K, D)
33 let idx: i64 = indices[k]
34 var j: i64 = 0
35 while j < D { out[j] = out[j] + scb[idx*D + j]; j = j + 1 }
36 k = k + 1
37 }
38 return 0
39}
40// squared-error distortion of v vs a reconstruction out[D]
41func rvq_dist(v: *i64, out: *i64, D: i64) -> i64 {
42 var s: i64 = 0
43 var i: i64 = 0
44 while i < D { let d: i64 = v[i] - out[i]; s = s + d*d; i = i + 1 }
45 return s
46}