nx_nofloat_autograd.nx source
↩ module page · 774 lines · 39001 B
1// nx_nofloat_autograd.nx -- NO-FLOAT (integer Q16 fixed-point) reverse-mode TENSOR autograd: the genuine
2// MISSING GENERATION in the no-float lineage (genealogy DeepMind ladder, operator first law = don't lie).
3//
4// HONEST landscape this fills (verified by reading the organs, 2026-06-22):
5// - GENERAL reverse-mode autograd ALREADY EXISTS, but on the SOFTWARE-FLOAT tower: nx_autograd /
6// nx_autograd_tensor / nx_tgrad_core all build on nx_f32 (IEEE-754 emulated). Float, not no-float.
7// - NO-FLOAT (integer Q16) training ALREADY EXISTS, but only nx_nn_train = a SINGLE linear layer with a
8// HAND-CODED analytic gradient. No general graph, no composition.
9// - What was genuinely missing = a GENERAL reverse-mode autograd in the SAME integer Q16 tower the
10// no-float Qwen inference pipeline actually uses (nx_nofloat_llm: Q16=65536, qmul=(a*b)>>16,
11// accumulate-then-shift matmul). This file is that bridge: a tape of arbitrary ops + ONE reverse
12// sweep, so an MLP (matvec->vadd->relu->matvec->vadd->mse) trains by the SAME loop that will train the
13// transformer -- all deterministic integer arithmetic (no nx_f32, no FPU, no GPU, no third-party AD).
14//
15// TAPE: stride-7 nodes {op, ai, bi, rows, cols, valp, gradp}; valp/gradp index a per-build BUMP ARENA of
16// Q16 cells (st[1] counter; no allocator -> deterministic). Forward EAGER (Q16), backward ONE reverse sweep
17// (nodes built in topological order, so reverse index order visits parents before children).
18// Backward identities (Q16) = the proven float identities with nx_f32_mul->qmul, nx_f32_add->integer +:
19// matvec y=W*x: dW[i][j] += qmul(gy[i], x[j]); dx[j] += qmul(W[i][j], gy[i])
20// vadd y=a+b: ga += gy ; gb += gy
21// relu: gx += gy iff the INPUT cell was > 0
22// mse L=(1/n)sum(p-t)^2: dp_i += qmul(gL, (2*(p_i-t_i))/n)
23// Determinism EXCEED-AXIS: integer add is EXACTLY associative -> bit-identical training independent of
24// order (the f32 tower is only bit-identical because it pins ONE order; here it is structural).
25// genealogy_id: linnainmaa_1970_reverse_mode_ad + rumelhart_1986_backprop, ported to fixed-point Q16
26// lineage_id: sovereign_nofloat_tensor_tape_autograd_v1
27// license_tier: ORIGINAL No `main` (pure library; proof lives in nx_nofloat_autograd_gate.nx).
28import "nx_syscalls.nx"
29const NFA_MAGIC_5040: i64 = 5040
30const NFA_MAGIC_362880: i64 = 362880
31const NFA_MAGIC_40320: i64 = 40320
32const NFA_MAGIC_45426: i64 = 45426
33
34const NFA_LEAF: i64 = 0
35const NFA_MATVEC: i64 = 1
36const NFA_VADD: i64 = 2
37const NFA_RELU: i64 = 3
38const NFA_MSE: i64 = 4
39const NFA_SOFTMAX: i64 = 5 // attention nonlinearity
40const NFA_SILU: i64 = 6 // FFN activation (x*sigmoid(x))
41const NFA_RMSNORM: i64 = 7 // transformer normalization
42const NFA_MATMUL: i64 = 8 // C[m,p] = A[m,k] . B[k,p]
43const NFA_MATMUL_NT: i64 = 9 // S[m,p] = A[m,k] . B[p,k]^T (the Q.K^T attention contraction)
44const NFA_CMUL: i64 = 10 // elementwise scale by a Q16 constant (bi holds the constant, not a node)
45const NFA_SOFTMAX_ROWS: i64 = 11 // per-row softmax (bi=1 -> CAUSAL: row i normalizes over j<=i only)
46const NFA_ROPE: i64 = 12 // rotary position embedding on [T,hd] (parameter-free; pos = row index)
47const NFA_HADAMARD: i64 = 13 // elementwise vector multiply a (*) b (the SwiGLU gate)
48const NFA_RMSNORM_ROWS: i64 = 14 // PER-ROW (per-token) RMSNorm over the cols of an [T,d] node
49const NFA_EMBED: i64 = 15 // gather rows of E[V,dm] by integer token ids -> X[T,dm] (bi = ids ptr)
50const NFA_SOFTCE_ROWS: i64 = 16 // fused per-row softmax cross-entropy over targets (bi = target ids ptr)
51const NFA_SLICE_COLS: i64 = 17 // extract cols [c0,c0+w) of an [T,C] node -> [T,w] (bi = c0); for multi-head
52const NFA_CONCAT_COLS: i64 = 18 // concat two [T,*] nodes column-wise -> [T, wa+wb]; for multi-head
53const NFA_Q16: i64 = 65536 // 1.0 in Q16 fixed-point (matches nx_nofloat_llm)
54// Q16 transcendental constants -- copied VERBATIM from the canonical nx_nofloat_llm so the TRAINING forward
55// uses the SAME fixed-point math as inference (DRY tension noted: a future shared nx_nofloat_math imported by
56// both is the clean resolution; kept inline so this lib still imports only nx_syscalls = no-float purity is
57// trivially auditable).
58const NFA_LOG2E: i64 = 94548
59const NFA_PC0: i64 = 65536
60const NFA_PC1: i64 = 45426
61const NFA_PC2: i64 = 15743
62const NFA_PC3: i64 = 4367
63// RoPE: Q16 angle constants + Qwen2.5 rope base ln(1e6) (verbatim from nx_nofloat_llm)
64const NFA_HALF_PI: i64 = 102944
65const NFA_PI: i64 = 205887
66const NFA_3HALF_PI: i64 = 308831
67const NFA_TWO_PI: i64 = 411775
68const NFA_LN_BASE: i64 = 905421
69
70// Q16 multiply: (a*b)>>16. Accumulate in full i64 precision, ONE shift -- the no-float matmul convention.
71func nfa_qmul(a: i64, b: i64) -> i64 { return (a * b) >> 16 }
72
73// ---- Q16 fixed-point math primitives (mirror nx_nofloat_llm; pure integer, no float) ----
74func nfa_isqrt(v: i64) -> i64 { if v<=0 { return 0 } if v<4 { return 1 } var x: i64=v; var y: i64=(x+1)>>1; var go: i64=1; while go==1 { if y<x { x=y; y=(x+v/x)>>1 } else { go=0 } } return x }
75// exp for arg <= 0 (softmax/sigmoid call it only after a max-subtract -> non-positive arg); returns Q16.
76func nfa_fxexp(x: i64) -> i64 { var xm: i64=0-x; if x>0 { xm=0 } let ym: i64=(xm*NFA_LOG2E)>>16; let yi: i64=ym>>16; let yf: i64=ym-(yi<<16); let g: i64=NFA_Q16-yf; var t: i64=NFA_PC3; t=NFA_PC2+((g*t)>>16); t=NFA_PC1+((g*t)>>16); t=NFA_PC0+((g*t)>>16); t=t>>1; if yi>=31 { return 0 } return t>>yi }
77func nfa_sigmoid(x: i64) -> i64 { if x>=0 { let ex: i64=nfa_fxexp(0-x); return (NFA_Q16*NFA_Q16)/(NFA_Q16+ex) } let ex: i64=nfa_fxexp(x); let sp: i64=(NFA_Q16*NFA_Q16)/(NFA_Q16+ex); return NFA_Q16-sp }
78func nfa_siluf(x: i64) -> i64 { return nfa_qmul(x, nfa_sigmoid(x)) } // silu(x) = x*sigmoid(x)
79func nfa_silud(x: i64) -> i64 { let s: i64=nfa_sigmoid(x); return s + nfa_qmul(x, nfa_qmul(s, NFA_Q16 - s)) } // silu'(x) = s + x*s*(1-s)
80// Q16 sin/cos (Taylor + quadrant reduction), verbatim from nx_nofloat_llm -- for RoPE rotation angles.
81func nfa_sinq(x: i64) -> i64 { let x2: i64=nfa_qmul(x,x); let x3: i64=nfa_qmul(x2,x); let x5: i64=nfa_qmul(x3,x2); let x7: i64=nfa_qmul(x5,x2); let x9: i64=nfa_qmul(x7,x2); return x - x3/6 + x5/120 - x7/NFA_MAGIC_5040 + x9/NFA_MAGIC_362880 }
82func nfa_cosq(x: i64) -> i64 { let x2: i64=nfa_qmul(x,x); let x4: i64=nfa_qmul(x2,x2); let x6: i64=nfa_qmul(x4,x2); let x8: i64=nfa_qmul(x6,x2); return NFA_Q16 - x2/2 + x4/24 - x6/720 + x8/NFA_MAGIC_40320 }
83func nfa_reduce2pi(a: i64) -> i64 { var t: i64=a; while t<0 { t=t+NFA_TWO_PI } while t>=NFA_TWO_PI { t=t-NFA_TWO_PI } return t }
84func nfa_sinf(a: i64) -> i64 { let t: i64=nfa_reduce2pi(a); if t<NFA_HALF_PI { return nfa_sinq(t) } if t<NFA_PI { return nfa_sinq(NFA_PI-t) } if t<NFA_3HALF_PI { return 0-nfa_sinq(t-NFA_PI) } return 0-nfa_sinq(NFA_TWO_PI-t) }
85func nfa_cosf(a: i64) -> i64 { let t: i64=nfa_reduce2pi(a); if t<NFA_HALF_PI { return nfa_cosq(t) } if t<NFA_PI { return 0-nfa_cosq(NFA_PI-t) } if t<NFA_3HALF_PI { return 0-nfa_cosq(t-NFA_PI) } return nfa_cosq(NFA_TWO_PI-t) }
86// Q16 natural log for x>0: x = m*2^e with m in [1,2); ln(x)=e*ln2 + 2*atanh((m-1)/(m+1)). The atanh series
87// converges fast because (m-1)/(m+1) in [0,1/3]. For the cross-entropy VALUE only (backward uses the exact
88// softmax identity, no log). 45426 = ln(2) in Q16.
89func nfa_ln(x: i64) -> i64 {
90 if x <= 0 { return 0 }
91 var e: i64 = 0; var m: i64 = x
92 while m >= 2*NFA_Q16 { m = m >> 1; e = e + 1 }
93 while m < NFA_Q16 { m = m << 1; e = e - 1 }
94 let y: i64 = ((m - NFA_Q16) << 16) / (m + NFA_Q16)
95 let y2: i64 = nfa_qmul(y, y)
96 var term: i64 = y; var sum: i64 = y
97 term = nfa_qmul(term, y2); sum = sum + term/3
98 term = nfa_qmul(term, y2); sum = sum + term/5
99 term = nfa_qmul(term, y2); sum = sum + term/7
100 term = nfa_qmul(term, y2); sum = sum + term/9
101 return e * NFA_MAGIC_45426 + 2*sum
102}
103
104// st[0] = next node index, st[1] = next free arena cell. Allocate a node, return its index.
105func nfa_new(tape: *i64, st: *i64, op: i64, ai: i64, bi: i64, rows: i64, cols: i64) -> i64 {
106 let k: i64 = st[0]
107 let off: i64 = st[1]
108 tape[7*k+0] = op; tape[7*k+1] = ai; tape[7*k+2] = bi
109 tape[7*k+3] = rows; tape[7*k+4] = cols
110 tape[7*k+5] = off; tape[7*k+6] = off
111 st[0] = k + 1
112 st[1] = off + rows * cols
113 return k
114}
115
116// leaf holding rows*cols Q16 cells copied from src[soff..].
117func nfa_leaf(tape: *i64, vals: *i64, st: *i64, rows: i64, cols: i64, src: *i64, soff: i64) -> i64 {
118 let k: i64 = nfa_new(tape, st, NFA_LEAF, 0 - 1, 0 - 1, rows, cols)
119 let off: i64 = tape[7*k+5]
120 let sz: i64 = rows * cols
121 var i: i64 = 0
122 while i < sz { vals[off + i] = src[soff + i]; i = i + 1 }
123 return k
124}
125
126// y = W * x (W is r*c, x is c*1, y is r*1): full-precision accumulate, ONE shift (no per-term underflow).
127func nfa_matvec(tape: *i64, vals: *i64, st: *i64, aW: i64, bx: i64) -> i64 {
128 let r: i64 = tape[7*aW+3]
129 let c: i64 = tape[7*aW+4]
130 let k: i64 = nfa_new(tape, st, NFA_MATVEC, aW, bx, r, 1)
131 let offy: i64 = tape[7*k+5]
132 let offW: i64 = tape[7*aW+5]
133 let offx: i64 = tape[7*bx+5]
134 var i: i64 = 0
135 while i < r {
136 var acc: i64 = 0
137 var j: i64 = 0
138 while j < c { acc = acc + vals[offW + i*c + j] * vals[offx + j]; j = j + 1 }
139 vals[offy + i] = acc >> 16
140 i = i + 1
141 }
142 return k
143}
144
145// y = a + b (same shape)
146func nfa_vadd(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
147 let r: i64 = tape[7*a+3]
148 let c: i64 = tape[7*a+4]
149 let k: i64 = nfa_new(tape, st, NFA_VADD, a, b, r, c)
150 let offy: i64 = tape[7*k+5]
151 let offa: i64 = tape[7*a+5]
152 let offb: i64 = tape[7*b+5]
153 let sz: i64 = r * c
154 var i: i64 = 0
155 while i < sz { vals[offy + i] = vals[offa + i] + vals[offb + i]; i = i + 1 }
156 return k
157}
158
159// y = relu(a) elementwise
160func nfa_relu(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
161 let r: i64 = tape[7*a+3]
162 let c: i64 = tape[7*a+4]
163 let k: i64 = nfa_new(tape, st, NFA_RELU, a, 0 - 1, r, c)
164 let offy: i64 = tape[7*k+5]
165 let offa: i64 = tape[7*a+5]
166 let sz: i64 = r * c
167 var i: i64 = 0
168 while i < sz {
169 var v: i64 = vals[offa + i]
170 if v < 0 { v = 0 }
171 vals[offy + i] = v
172 i = i + 1
173 }
174 return k
175}
176
177// L = (1/n) sum_i (pred_i - target_i)^2 (scalar, 1 cell). Q16: di=p_q-t_q is Q16, di*di is Q32,
178// L_q = (sum di*di) / (Q16 * n) -> back to Q16 of the mean squared error.
179func nfa_mse(tape: *i64, vals: *i64, st: *i64, pred: i64, target: i64) -> i64 {
180 let np: i64 = tape[7*pred+3] * tape[7*pred+4]
181 let k: i64 = nfa_new(tape, st, NFA_MSE, pred, target, 1, 1)
182 let offL: i64 = tape[7*k+5]
183 let offp: i64 = tape[7*pred+5]
184 let offt: i64 = tape[7*target+5]
185 var acc: i64 = 0
186 var i: i64 = 0
187 while i < np {
188 let di: i64 = vals[offp + i] - vals[offt + i]
189 acc = acc + di * di
190 i = i + 1
191 }
192 vals[offL] = acc / (NFA_Q16 * np)
193 return k
194}
195
196// ---- transformer-defining ops (the Qwen sublayer nonlinearities), each with a Q16 reverse-mode backward ----
197// y = softmax(a) over the node's n=rows*cols cells (one vector). Stores y for the backward Jacobian-vector product.
198func nfa_softmax(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
199 let r: i64 = tape[7*a+3]; let c: i64 = tape[7*a+4]; let n: i64 = r*c
200 let k: i64 = nfa_new(tape, st, NFA_SOFTMAX, a, 0 - 1, r, c)
201 let offy: i64 = tape[7*k+5]; let offa: i64 = tape[7*a+5]
202 var mx: i64 = vals[offa]; var i: i64 = 1
203 while i < n { if vals[offa+i] > mx { mx = vals[offa+i] } i = i + 1 }
204 var sum: i64 = 0; i = 0
205 while i < n { let e: i64 = nfa_fxexp(vals[offa+i] - mx); vals[offy+i] = e; sum = sum + e; i = i + 1 }
206 if sum <= 0 { sum = 1 }
207 i = 0
208 while i < n { vals[offy+i] = (vals[offy+i] << 16) / sum; i = i + 1 }
209 return k
210}
211// y = silu(a) elementwise = a*sigmoid(a) (the FFN activation)
212func nfa_silu(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
213 let r: i64 = tape[7*a+3]; let c: i64 = tape[7*a+4]; let n: i64 = r*c
214 let k: i64 = nfa_new(tape, st, NFA_SILU, a, 0 - 1, r, c)
215 let offy: i64 = tape[7*k+5]; let offa: i64 = tape[7*a+5]
216 var i: i64 = 0
217 while i < n { vals[offy+i] = nfa_siluf(vals[offa+i]); i = i + 1 }
218 return k
219}
220// y = rmsnorm(a) (no gamma): y_i = a_i / sqrt(mean(a^2)). The transformer normalization.
221func nfa_rmsnorm(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
222 let r: i64 = tape[7*a+3]; let c: i64 = tape[7*a+4]; let n: i64 = r*c
223 let k: i64 = nfa_new(tape, st, NFA_RMSNORM, a, 0 - 1, r, c)
224 let offy: i64 = tape[7*k+5]; let offa: i64 = tape[7*a+5]
225 var ss: i64 = 0; var i: i64 = 0
226 while i < n { ss = ss + nfa_qmul(vals[offa+i], vals[offa+i]); i = i + 1 }
227 let ms: i64 = ss / n
228 let sd: i64 = nfa_isqrt((ms + 1) << 16)
229 i = 0
230 while i < n { if sd > 0 { vals[offy+i] = (vals[offa+i] << 16) / sd } else { vals[offy+i] = 0 } i = i + 1 }
231 return k
232}
233
234// ---- attention-contraction ops (the QK^T / AV matmuls + causal row-softmax), each with Q16 backward ----
235// C = A . B (A is [m,k], B is [k,p] -> C is [m,p]): full-precision accumulate, ONE shift.
236func nfa_matmul(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
237 let m: i64 = tape[7*a+3]; let kk: i64 = tape[7*a+4]; let p: i64 = tape[7*b+4]
238 let nd: i64 = nfa_new(tape, st, NFA_MATMUL, a, b, m, p)
239 let offC: i64 = tape[7*nd+5]; let offA: i64 = tape[7*a+5]; let offB: i64 = tape[7*b+5]
240 var i: i64 = 0
241 while i < m {
242 var j: i64 = 0
243 while j < p {
244 var acc: i64 = 0; var l: i64 = 0
245 while l < kk { acc = acc + vals[offA + i*kk + l] * vals[offB + l*p + j]; l = l + 1 }
246 vals[offC + i*p + j] = acc >> 16
247 j = j + 1
248 }
249 i = i + 1
250 }
251 return nd
252}
253// S = A . B^T (A is [m,k], B is [p,k] -> S is [m,p]): S[i][j] = sum_l A[i][l] B[j][l] (the Q.K^T form).
254func nfa_matmul_nt(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
255 let m: i64 = tape[7*a+3]; let kk: i64 = tape[7*a+4]; let p: i64 = tape[7*b+3]
256 let nd: i64 = nfa_new(tape, st, NFA_MATMUL_NT, a, b, m, p)
257 let offS: i64 = tape[7*nd+5]; let offA: i64 = tape[7*a+5]; let offB: i64 = tape[7*b+5]
258 var i: i64 = 0
259 while i < m {
260 var j: i64 = 0
261 while j < p {
262 var acc: i64 = 0; var l: i64 = 0
263 while l < kk { acc = acc + vals[offA + i*kk + l] * vals[offB + j*kk + l]; l = l + 1 }
264 vals[offS + i*p + j] = acc >> 16
265 j = j + 1
266 }
267 i = i + 1
268 }
269 return nd
270}
271// y = a * c elementwise, c a Q16 constant (stored in bi). The 1/sqrt(d) attention scale.
272func nfa_cmul(tape: *i64, vals: *i64, st: *i64, a: i64, c_q: i64) -> i64 {
273 let r: i64 = tape[7*a+3]; let c: i64 = tape[7*a+4]; let n: i64 = r*c
274 let nd: i64 = nfa_new(tape, st, NFA_CMUL, a, c_q, r, c)
275 let offy: i64 = tape[7*nd+5]; let offa: i64 = tape[7*a+5]
276 var i: i64 = 0
277 while i < n { vals[offy+i] = nfa_qmul(vals[offa+i], c_q); i = i + 1 }
278 return nd
279}
280// y = per-row softmax(a); causal=1 -> row i normalizes over columns j<=i only (the rest are 0).
281func nfa_softmax_rows(tape: *i64, vals: *i64, st: *i64, a: i64, causal: i64) -> i64 {
282 let r: i64 = tape[7*a+3]; let c: i64 = tape[7*a+4]
283 let nd: i64 = nfa_new(tape, st, NFA_SOFTMAX_ROWS, a, causal, r, c)
284 let offy: i64 = tape[7*nd+5]; let offa: i64 = tape[7*a+5]
285 var i: i64 = 0
286 while i < r {
287 var lim: i64 = c
288 if causal == 1 { lim = i + 1 }
289 let base: i64 = i * c
290 var mx: i64 = vals[offa+base]; var j: i64 = 1
291 while j < lim { if vals[offa+base+j] > mx { mx = vals[offa+base+j] } j = j + 1 }
292 var sum: i64 = 0; j = 0
293 while j < lim { let e: i64 = nfa_fxexp(vals[offa+base+j] - mx); vals[offy+base+j] = e; sum = sum + e; j = j + 1 }
294 j = lim
295 while j < c { vals[offy+base+j] = 0; j = j + 1 }
296 if sum <= 0 { sum = 1 }
297 j = 0
298 while j < lim { vals[offy+base+j] = (vals[offy+base+j] << 16) / sum; j = j + 1 }
299 i = i + 1
300 }
301 return nd
302}
303// y = RoPE(a): rotary position embedding on [T, hd] (hd even). Row t (= position t) rotates each (2i,2i+1)
304// pair by angle t*theta_i, theta_i = base^(-2i/hd). Parameter-free; backward = rotate by -angle (orthogonal).
305func nfa_rope(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
306 let T: i64 = tape[7*a+3]; let hd: i64 = tape[7*a+4]; let np: i64 = hd/2
307 let nd: i64 = nfa_new(tape, st, NFA_ROPE, a, 0 - 1, T, hd)
308 let offy: i64 = tape[7*nd+5]; let offa: i64 = tape[7*a+5]
309 var t: i64 = 0
310 while t < T {
311 var i: i64 = 0
312 while i < np {
313 let theta: i64 = nfa_fxexp(0 - (i*NFA_LN_BASE)/np)
314 let ang: i64 = t * theta
315 let c: i64 = nfa_cosf(ang); let s: i64 = nfa_sinf(ang)
316 let av: i64 = vals[offa + t*hd + 2*i]; let bv: i64 = vals[offa + t*hd + 2*i + 1]
317 vals[offy + t*hd + 2*i] = nfa_qmul(av,c) - nfa_qmul(bv,s)
318 vals[offy + t*hd + 2*i + 1] = nfa_qmul(av,s) + nfa_qmul(bv,c)
319 i = i + 1
320 }
321 t = t + 1
322 }
323 return nd
324}
325// y = a (*) b elementwise (Hadamard); the SwiGLU gate: silu(gate) (*) up.
326func nfa_hadamard(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
327 let r: i64 = tape[7*a+3]; let c: i64 = tape[7*a+4]; let n: i64 = r*c
328 let nd: i64 = nfa_new(tape, st, NFA_HADAMARD, a, b, r, c)
329 let offy: i64 = tape[7*nd+5]; let offa: i64 = tape[7*a+5]; let offb: i64 = tape[7*b+5]
330 var i: i64 = 0
331 while i < n { vals[offy+i] = nfa_qmul(vals[offa+i], vals[offb+i]); i = i + 1 }
332 return nd
333}
334// y = rmsnorm of EACH ROW (per-token) of an [r,c] node: y[i][j] = x[i][j]/sqrt(mean_j(x[i]^2)). Transformer pre-norm.
335func nfa_rmsnorm_rows(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
336 let r: i64 = tape[7*a+3]; let c: i64 = tape[7*a+4]
337 let nd: i64 = nfa_new(tape, st, NFA_RMSNORM_ROWS, a, 0 - 1, r, c)
338 let offy: i64 = tape[7*nd+5]; let offa: i64 = tape[7*a+5]
339 var i: i64 = 0
340 while i < r {
341 let base: i64 = i * c
342 var ss: i64 = 0; var j: i64 = 0
343 while j < c { ss = ss + nfa_qmul(vals[offa+base+j], vals[offa+base+j]); j = j + 1 }
344 let ms: i64 = ss / c
345 let sd: i64 = nfa_isqrt((ms + 1) << 16)
346 j = 0
347 while j < c { if sd > 0 { vals[offy+base+j] = (vals[offa+base+j] << 16) / sd } else { vals[offy+base+j] = 0 } j = j + 1 }
348 i = i + 1
349 }
350 return nd
351}
352// ---- LM head ops: embedding gather + fused softmax cross-entropy (ids are RAW ints, not Q16; no grad to ids) ----
353// X = embed(E, ids): E is [V, dm], ids is *i64 of T token ids -> X[T, dm], X[t] = E[ids[t]]. ids ptr stashed in bi.
354func nfa_embed(tape: *i64, vals: *i64, st: *i64, e_node: i64, ids: *i64, T: i64) -> i64 {
355 let dm: i64 = tape[7*e_node+4]
356 let nd: i64 = nfa_new(tape, st, NFA_EMBED, e_node, ids as i64, T, dm)
357 let offx: i64 = tape[7*nd+5]; let offE: i64 = tape[7*e_node+5]
358 var t: i64 = 0
359 while t < T {
360 let id: i64 = ids[t]
361 var j: i64 = 0
362 while j < dm { vals[offx + t*dm + j] = vals[offE + id*dm + j]; j = j + 1 }
363 t = t + 1
364 }
365 return nd
366}
367// L = (1/T) sum_t [ logsumexp(logits[t]) - logits[t][id_t] ]; logits is [T,V], tgt is *i64 of T target ids.
368// Stored value = the scalar Q16 loss. Backward (in nfa_backward) uses the exact identity dlogit=(softmax-onehot)/T.
369func nfa_softce_rows(tape: *i64, vals: *i64, st: *i64, logits: i64, tgt: *i64) -> i64 {
370 let T: i64 = tape[7*logits+3]; let V: i64 = tape[7*logits+4]
371 let nd: i64 = nfa_new(tape, st, NFA_SOFTCE_ROWS, logits, tgt as i64, 1, 1)
372 let offL: i64 = tape[7*nd+5]; let offp: i64 = tape[7*logits+5]
373 var total: i64 = 0; var t: i64 = 0
374 while t < T {
375 let base: i64 = t*V
376 var mx: i64 = vals[offp+base]; var j: i64 = 1
377 while j < V { if vals[offp+base+j] > mx { mx = vals[offp+base+j] } j = j + 1 }
378 var sum: i64 = 0; j = 0
379 while j < V { sum = sum + nfa_fxexp(vals[offp+base+j] - mx); j = j + 1 }
380 let lse: i64 = mx + nfa_ln(sum)
381 total = total + (lse - vals[offp+base+tgt[t]])
382 t = t + 1
383 }
384 vals[offL] = total / T
385 return nd
386}
387// ---- column slice / concat (the head-split bookkeeping for MULTI-HEAD attention) ----
388// y = a[:, c0:c0+w] (a is [T,C] -> y is [T,w]); c0 stored in bi. backward scatters gy into a's columns.
389func nfa_slice_cols(tape: *i64, vals: *i64, st: *i64, a: i64, c0: i64, w: i64) -> i64 {
390 let T: i64 = tape[7*a+3]; let C: i64 = tape[7*a+4]
391 let nd: i64 = nfa_new(tape, st, NFA_SLICE_COLS, a, c0, T, w)
392 let offy: i64 = tape[7*nd+5]; let offa: i64 = tape[7*a+5]
393 var t: i64 = 0
394 while t < T { var j: i64 = 0; while j < w { vals[offy + t*w + j] = vals[offa + t*C + c0 + j]; j = j + 1 } t = t + 1 }
395 return nd
396}
397// y = [a | b] (a is [T,wa], b is [T,wb] -> y is [T,wa+wb]). backward splits gy back to a and b.
398func nfa_concat_cols(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
399 let T: i64 = tape[7*a+3]; let wa: i64 = tape[7*a+4]; let wb: i64 = tape[7*b+4]; let wc: i64 = wa + wb
400 let nd: i64 = nfa_new(tape, st, NFA_CONCAT_COLS, a, b, T, wc)
401 let offy: i64 = tape[7*nd+5]; let offa: i64 = tape[7*a+5]; let offb: i64 = tape[7*b+5]
402 var t: i64 = 0
403 while t < T {
404 var j: i64 = 0
405 while j < wa { vals[offy + t*wc + j] = vals[offa + t*wa + j]; j = j + 1 }
406 j = 0
407 while j < wb { vals[offy + t*wc + wa + j] = vals[offb + t*wb + j]; j = j + 1 }
408 t = t + 1
409 }
410 return nd
411}
412
413func nfa_val(tape: *i64, vals: *i64, k: i64, c: i64) -> i64 { return vals[tape[7*k+5] + c] }
414func nfa_grad(tape: *i64, grads: *i64, k: i64, c: i64) -> i64 { return grads[tape[7*k+6] + c] }
415
416// reverse sweep: zero all grads, seed grad[root]=Q16(1.0), accumulate the identities backward.
417func nfa_backward(tape: *i64, vals: *i64, grads: *i64, n: i64, root: i64) -> i64 {
418 var k: i64 = 0
419 while k < n {
420 let off: i64 = tape[7*k+5]
421 let sz: i64 = tape[7*k+3] * tape[7*k+4]
422 var i: i64 = 0
423 while i < sz { grads[off + i] = 0; i = i + 1 }
424 k = k + 1
425 }
426 grads[tape[7*root+6] + 0] = NFA_Q16
427 k = n - 1
428 while k >= 0 {
429 let op: i64 = tape[7*k+0]
430 if op == NFA_MATVEC {
431 let aW: i64 = tape[7*k+1]
432 let bx: i64 = tape[7*k+2]
433 let r: i64 = tape[7*k+3]
434 let c: i64 = tape[7*aW+4]
435 let offgy: i64 = tape[7*k+6]
436 let offW: i64 = tape[7*aW+5]
437 let gW: i64 = tape[7*aW+6]
438 let offx: i64 = tape[7*bx+5]
439 let gx: i64 = tape[7*bx+6]
440 var i: i64 = 0
441 while i < r {
442 let gyi: i64 = grads[offgy + i]
443 var j: i64 = 0
444 while j < c {
445 grads[gW + i*c + j] = grads[gW + i*c + j] + nfa_qmul(gyi, vals[offx + j])
446 grads[gx + j] = grads[gx + j] + nfa_qmul(vals[offW + i*c + j], gyi)
447 j = j + 1
448 }
449 i = i + 1
450 }
451 }
452 if op == NFA_VADD {
453 let a: i64 = tape[7*k+1]
454 let b: i64 = tape[7*k+2]
455 let sz: i64 = tape[7*k+3] * tape[7*k+4]
456 let offgy: i64 = tape[7*k+6]
457 let ga: i64 = tape[7*a+6]
458 let gb: i64 = tape[7*b+6]
459 var i: i64 = 0
460 while i < sz {
461 let g: i64 = grads[offgy + i]
462 grads[ga + i] = grads[ga + i] + g
463 grads[gb + i] = grads[gb + i] + g
464 i = i + 1
465 }
466 }
467 if op == NFA_RELU {
468 let a: i64 = tape[7*k+1]
469 let sz: i64 = tape[7*k+3] * tape[7*k+4]
470 let offgy: i64 = tape[7*k+6]
471 let ga: i64 = tape[7*a+6]
472 let offa: i64 = tape[7*a+5]
473 var i: i64 = 0
474 while i < sz {
475 if vals[offa + i] > 0 { grads[ga + i] = grads[ga + i] + grads[offgy + i] }
476 i = i + 1
477 }
478 }
479 if op == NFA_MSE {
480 let a: i64 = tape[7*k+1]
481 let b: i64 = tape[7*k+2]
482 let np: i64 = tape[7*a+3] * tape[7*a+4]
483 let offp: i64 = tape[7*a+5]
484 let gp: i64 = tape[7*a+6]
485 let offt: i64 = tape[7*b+5]
486 let gL: i64 = grads[tape[7*k+6] + 0]
487 var i: i64 = 0
488 while i < np {
489 let di: i64 = vals[offp + i] - vals[offt + i]
490 grads[gp + i] = grads[gp + i] + nfa_qmul(gL, (2 * di) / np)
491 i = i + 1
492 }
493 }
494 if op == NFA_SOFTMAX {
495 // dx_i = y_i * (gy_i - sum_j y_j gy_j) -- the softmax Jacobian-vector product
496 let a: i64 = tape[7*k+1]
497 let n: i64 = tape[7*k+3] * tape[7*k+4]
498 let offy: i64 = tape[7*k+5]
499 let offgy: i64 = tape[7*k+6]
500 let ga: i64 = tape[7*a+6]
501 var dot: i64 = 0; var j: i64 = 0
502 while j < n { dot = dot + nfa_qmul(vals[offy+j], grads[offgy+j]); j = j + 1 }
503 var i: i64 = 0
504 while i < n { grads[ga+i] = grads[ga+i] + nfa_qmul(vals[offy+i], grads[offgy+i] - dot); i = i + 1 }
505 }
506 if op == NFA_SILU {
507 // dx_i = gy_i * silu'(x_i)
508 let a: i64 = tape[7*k+1]
509 let n: i64 = tape[7*k+3] * tape[7*k+4]
510 let offgy: i64 = tape[7*k+6]
511 let offa: i64 = tape[7*a+5]
512 let ga: i64 = tape[7*a+6]
513 var i: i64 = 0
514 while i < n { grads[ga+i] = grads[ga+i] + nfa_qmul(grads[offgy+i], nfa_silud(vals[offa+i])); i = i + 1 }
515 }
516 if op == NFA_RMSNORM {
517 // dx_j = gy_j/r - y_j*dot/(n*ms), r=rms, ms=mean(x^2), dot=sum_i gy_i x_i (Q16; recompute fwd stats)
518 let a: i64 = tape[7*k+1]
519 let n: i64 = tape[7*k+3] * tape[7*k+4]
520 let offy: i64 = tape[7*k+5]
521 let offgy: i64 = tape[7*k+6]
522 let offa: i64 = tape[7*a+5]
523 let ga: i64 = tape[7*a+6]
524 var ss: i64 = 0; var i: i64 = 0
525 while i < n { ss = ss + nfa_qmul(vals[offa+i], vals[offa+i]); i = i + 1 }
526 let ms: i64 = ss / n
527 let sd: i64 = nfa_isqrt((ms + 1) << 16)
528 var dot: i64 = 0; i = 0
529 while i < n { dot = dot + nfa_qmul(grads[offgy+i], vals[offa+i]); i = i + 1 }
530 if sd > 0 { if ms > 0 {
531 i = 0
532 while i < n {
533 let t1: i64 = (grads[offgy+i] << 16) / sd
534 let t2: i64 = (vals[offy+i] * dot) / (n * ms)
535 grads[ga+i] = grads[ga+i] + t1 - t2
536 i = i + 1
537 }
538 } }
539 }
540 if op == NFA_MATMUL {
541 // C=A.B: dA[i][l] += sum_j gy[i][j] B[l][j]; dB[l][j] += sum_i A[i][l] gy[i][j]
542 let a: i64 = tape[7*k+1]; let b: i64 = tape[7*k+2]
543 let m: i64 = tape[7*k+3]; let p: i64 = tape[7*k+4]; let kk: i64 = tape[7*a+4]
544 let offgy: i64 = tape[7*k+6]; let offA: i64 = tape[7*a+5]; let gA: i64 = tape[7*a+6]; let offB: i64 = tape[7*b+5]; let gB: i64 = tape[7*b+6]
545 var i: i64 = 0
546 while i < m {
547 var j: i64 = 0
548 while j < p {
549 let g: i64 = grads[offgy + i*p + j]; var l: i64 = 0
550 while l < kk {
551 grads[gA + i*kk + l] = grads[gA + i*kk + l] + nfa_qmul(g, vals[offB + l*p + j])
552 grads[gB + l*p + j] = grads[gB + l*p + j] + nfa_qmul(vals[offA + i*kk + l], g)
553 l = l + 1
554 }
555 j = j + 1
556 }
557 i = i + 1
558 }
559 }
560 if op == NFA_MATMUL_NT {
561 // S=A.B^T: dA[i][l] += sum_j gy[i][j] B[j][l]; dB[j][l] += sum_i gy[i][j] A[i][l]
562 let a: i64 = tape[7*k+1]; let b: i64 = tape[7*k+2]
563 let m: i64 = tape[7*k+3]; let p: i64 = tape[7*k+4]; let kk: i64 = tape[7*a+4]
564 let offgy: i64 = tape[7*k+6]; let offA: i64 = tape[7*a+5]; let gA: i64 = tape[7*a+6]; let offB: i64 = tape[7*b+5]; let gB: i64 = tape[7*b+6]
565 var i: i64 = 0
566 while i < m {
567 var j: i64 = 0
568 while j < p {
569 let g: i64 = grads[offgy + i*p + j]; var l: i64 = 0
570 while l < kk {
571 grads[gA + i*kk + l] = grads[gA + i*kk + l] + nfa_qmul(g, vals[offB + j*kk + l])
572 grads[gB + j*kk + l] = grads[gB + j*kk + l] + nfa_qmul(g, vals[offA + i*kk + l])
573 l = l + 1
574 }
575 j = j + 1
576 }
577 i = i + 1
578 }
579 }
580 if op == NFA_CMUL {
581 let a: i64 = tape[7*k+1]; let c_q: i64 = tape[7*k+2]
582 let n: i64 = tape[7*k+3] * tape[7*k+4]
583 let offgy: i64 = tape[7*k+6]; let ga: i64 = tape[7*a+6]
584 var i: i64 = 0
585 while i < n { grads[ga+i] = grads[ga+i] + nfa_qmul(grads[offgy+i], c_q); i = i + 1 }
586 }
587 if op == NFA_SOFTMAX_ROWS {
588 // per-row Jacobian-vector product over j<lim (causal: lim=i+1); masked entries had 0 effect -> 0 grad
589 let a: i64 = tape[7*k+1]; let causal: i64 = tape[7*k+2]
590 let r: i64 = tape[7*k+3]; let c: i64 = tape[7*k+4]
591 let offy: i64 = tape[7*k+5]; let offgy: i64 = tape[7*k+6]; let ga: i64 = tape[7*a+6]
592 var i: i64 = 0
593 while i < r {
594 var lim: i64 = c
595 if causal == 1 { lim = i + 1 }
596 let base: i64 = i * c
597 var dot: i64 = 0; var j: i64 = 0
598 while j < lim { dot = dot + nfa_qmul(vals[offy+base+j], grads[offgy+base+j]); j = j + 1 }
599 j = 0
600 while j < lim { grads[ga+base+j] = grads[ga+base+j] + nfa_qmul(vals[offy+base+j], grads[offgy+base+j] - dot); j = j + 1 }
601 i = i + 1
602 }
603 }
604 if op == NFA_ROPE {
605 // backward = rotate gy by -ang (orthogonal Jacobian): gv[2i]=g0*c+g1*s; gv[2i+1]=g1*c-g0*s
606 let a: i64 = tape[7*k+1]; let T: i64 = tape[7*k+3]; let hd: i64 = tape[7*k+4]; let np: i64 = hd/2
607 let offgy: i64 = tape[7*k+6]; let ga: i64 = tape[7*a+6]
608 var t: i64 = 0
609 while t < T {
610 var i: i64 = 0
611 while i < np {
612 let theta: i64 = nfa_fxexp(0 - (i*NFA_LN_BASE)/np)
613 let ang: i64 = t * theta
614 let c: i64 = nfa_cosf(ang); let s: i64 = nfa_sinf(ang)
615 let g0: i64 = grads[offgy + t*hd + 2*i]; let g1: i64 = grads[offgy + t*hd + 2*i + 1]
616 grads[ga + t*hd + 2*i] = grads[ga + t*hd + 2*i] + nfa_qmul(g0,c) + nfa_qmul(g1,s)
617 grads[ga + t*hd + 2*i + 1] = grads[ga + t*hd + 2*i + 1] + nfa_qmul(g1,c) - nfa_qmul(g0,s)
618 i = i + 1
619 }
620 t = t + 1
621 }
622 }
623 if op == NFA_HADAMARD {
624 let a: i64 = tape[7*k+1]; let b: i64 = tape[7*k+2]
625 let n: i64 = tape[7*k+3] * tape[7*k+4]
626 let offgy: i64 = tape[7*k+6]; let offa: i64 = tape[7*a+5]; let offb: i64 = tape[7*b+5]
627 let ga: i64 = tape[7*a+6]; let gb: i64 = tape[7*b+6]
628 var i: i64 = 0
629 while i < n {
630 let g: i64 = grads[offgy+i]
631 grads[ga+i] = grads[ga+i] + nfa_qmul(g, vals[offb+i])
632 grads[gb+i] = grads[gb+i] + nfa_qmul(g, vals[offa+i])
633 i = i + 1
634 }
635 }
636 if op == NFA_RMSNORM_ROWS {
637 // per-row RMSNorm backward: dx_j = gy_j/r - y_j*dot/(c*ms) per row (dot = sum_j gy_j x_j), Q16
638 let a: i64 = tape[7*k+1]; let r: i64 = tape[7*k+3]; let c: i64 = tape[7*k+4]
639 let offy: i64 = tape[7*k+5]; let offgy: i64 = tape[7*k+6]; let offa: i64 = tape[7*a+5]; let ga: i64 = tape[7*a+6]
640 var i: i64 = 0
641 while i < r {
642 let base: i64 = i * c
643 var ss: i64 = 0; var j: i64 = 0
644 while j < c { ss = ss + nfa_qmul(vals[offa+base+j], vals[offa+base+j]); j = j + 1 }
645 let ms: i64 = ss / c
646 let sd: i64 = nfa_isqrt((ms + 1) << 16)
647 var dot: i64 = 0; j = 0
648 while j < c { dot = dot + nfa_qmul(grads[offgy+base+j], vals[offa+base+j]); j = j + 1 }
649 if sd > 0 { if ms > 0 {
650 j = 0
651 while j < c {
652 let t1: i64 = (grads[offgy+base+j] << 16) / sd
653 let t2: i64 = (vals[offy+base+j] * dot) / (c * ms)
654 grads[ga+base+j] = grads[ga+base+j] + t1 - t2
655 j = j + 1
656 }
657 } }
658 i = i + 1
659 }
660 }
661 if op == NFA_EMBED {
662 // scatter-ADD gy rows into dE[ids[t]] (a repeated token id accumulates grads); no grad to ids
663 let e_node: i64 = tape[7*k+1]; let ids: *i64 = tape[7*k+2] as *i64
664 let T: i64 = tape[7*k+3]; let dm: i64 = tape[7*k+4]
665 let offgy: i64 = tape[7*k+6]; let gE: i64 = tape[7*e_node+6]
666 var t: i64 = 0
667 while t < T {
668 let id: i64 = ids[t]
669 var j: i64 = 0
670 while j < dm { grads[gE + id*dm + j] = grads[gE + id*dm + j] + grads[offgy + t*dm + j]; j = j + 1 }
671 t = t + 1
672 }
673 }
674 if op == NFA_SOFTCE_ROWS {
675 // dlogit[t][j] = qmul(gL, (softmax[t][j] - onehot[t][j]) / T) -- the exact fused CE identity (no log)
676 let logits: i64 = tape[7*k+1]; let tgt: *i64 = tape[7*k+2] as *i64
677 let T: i64 = tape[7*logits+3]; let V: i64 = tape[7*logits+4]
678 let offp: i64 = tape[7*logits+5]; let gp: i64 = tape[7*logits+6]
679 let gL: i64 = grads[tape[7*k+6] + 0]
680 var t: i64 = 0
681 while t < T {
682 let base: i64 = t*V
683 var mx: i64 = vals[offp+base]; var j: i64 = 1
684 while j < V { if vals[offp+base+j] > mx { mx = vals[offp+base+j] } j = j + 1 }
685 var sum: i64 = 0; j = 0
686 while j < V { sum = sum + nfa_fxexp(vals[offp+base+j] - mx); j = j + 1 }
687 if sum <= 0 { sum = 1 }
688 let id: i64 = tgt[t]
689 j = 0
690 while j < V {
691 let sm: i64 = (nfa_fxexp(vals[offp+base+j] - mx) << 16) / sum
692 var oh: i64 = 0
693 if j == id { oh = NFA_Q16 }
694 grads[gp+base+j] = grads[gp+base+j] + nfa_qmul(gL, (sm - oh) / T)
695 j = j + 1
696 }
697 t = t + 1
698 }
699 }
700 if op == NFA_SLICE_COLS {
701 let a: i64 = tape[7*k+1]; let c0: i64 = tape[7*k+2]; let T: i64 = tape[7*k+3]; let w: i64 = tape[7*k+4]; let C: i64 = tape[7*a+4]
702 let offgy: i64 = tape[7*k+6]; let ga: i64 = tape[7*a+6]
703 var t: i64 = 0
704 while t < T { var j: i64 = 0; while j < w { grads[ga + t*C + c0 + j] = grads[ga + t*C + c0 + j] + grads[offgy + t*w + j]; j = j + 1 } t = t + 1 }
705 }
706 if op == NFA_CONCAT_COLS {
707 let a: i64 = tape[7*k+1]; let b: i64 = tape[7*k+2]; let T: i64 = tape[7*k+3]; let wc: i64 = tape[7*k+4]
708 let wa: i64 = tape[7*a+4]; let wb: i64 = tape[7*b+4]
709 let offgy: i64 = tape[7*k+6]; let ga: i64 = tape[7*a+6]; let gb: i64 = tape[7*b+6]
710 var t: i64 = 0
711 while t < T {
712 var j: i64 = 0
713 while j < wa { grads[ga + t*wa + j] = grads[ga + t*wa + j] + grads[offgy + t*wc + j]; j = j + 1 }
714 j = 0
715 while j < wb { grads[gb + t*wb + j] = grads[gb + t*wb + j] + grads[offgy + t*wc + wa + j]; j = j + 1 }
716 t = t + 1
717 }
718 }
719 k = k - 1
720 }
721 return 0
722}
723
724// SGD optimizer step on a parameter vector: w[i] -= lr * g[i], lr in Q16. Step = (lr_q*g_q)/Q16 with
725// integer division truncating TOWARD ZERO -> sign-symmetric (no arithmetic-shift floor bias). General
726// (any param vector, any Q16 rate); the same step that scales from this 2-weight demo to the full model.
727func nfa_sgd(w: *i64, g: *i64, n: i64, lr_q: i64) -> i64 {
728 var i: i64 = 0
729 while i < n { w[i] = w[i] - (lr_q * g[i]) / NFA_Q16; i = i + 1 }
730 return 0
731}
732
733// ERROR-FEEDBACK SGD: like nfa_sgd but sub-quantum step remainders ACCUMULATE in resid[n] (persistent,
734// zero-init) instead of being truncated away -- no gradient is ever lost to the /Q16 quantization. This is
735// the integer-training lesson the neural-image arc proved (small systematic gradients otherwise quantize to
736// ZERO updates, and models freeze in degenerate states, e.g. the reader's uniform-collapse absorbing state).
737// acc = resid + lr*g (raw Q32-ish); apply the whole-quantum part; keep the remainder (sign-symmetric).
738func nfa_sgd_ef(w: *i64, g: *i64, resid: *i64, n: i64, lr_q: i64) -> i64 {
739 var i: i64 = 0
740 while i < n {
741 let acc: i64 = resid[i] + lr_q * g[i]
742 let step: i64 = acc / NFA_Q16
743 w[i] = w[i] - step
744 resid[i] = acc - step * NFA_Q16
745 i = i + 1
746 }
747 return 0
748}
749
750// AdamW optimizer step (Q16 integer; port of nx_tgrad_core's ad_step to fixed-point). Per-param adaptive
751// moments + bias-correction -> normalizes each parameter's step by its own gradient scale, so it converges
752// where plain SGD (one global rate) struggles on ill-conditioned losses. m/v are persistent Q16 state (zero-init);
753// t = 1-based step. All integer -> bit-exact/deterministic. sqrt via nfa_isqrt. b1/b2/eps/wd/lr all Q16.
754func nfa_adamw(w: *i64, g: *i64, m: *i64, v: *i64, n: i64, lr: i64, b1: i64, b2: i64, eps: i64, wd: i64, t: i64) -> i64 {
755 var c1: i64 = NFA_Q16; var c2: i64 = NFA_Q16
756 var k: i64 = 0
757 while k < t { c1 = nfa_qmul(c1, b1); c2 = nfa_qmul(c2, b2); k = k + 1 } // b1^t, b2^t
758 let bc1: i64 = NFA_Q16 - c1
759 let bc2: i64 = NFA_Q16 - c2
760 var i: i64 = 0
761 while i < n {
762 m[i] = nfa_qmul(b1, m[i]) + nfa_qmul(NFA_Q16 - b1, g[i])
763 v[i] = nfa_qmul(b2, v[i]) + nfa_qmul(NFA_Q16 - b2, nfa_qmul(g[i], g[i]))
764 var mh: i64 = m[i]; if bc1 > 0 { mh = (m[i] << 16) / bc1 }
765 var vh: i64 = v[i]; if bc2 > 0 { vh = (v[i] << 16) / bc2 }
766 if vh < 0 { vh = 0 }
767 let sq: i64 = nfa_isqrt(vh << 16) // sqrt(vh) in Q16
768 let denom: i64 = sq + eps
769 var upd: i64 = 0; if denom > 0 { upd = (mh << 16) / denom } // mhat/(sqrt(vhat)+eps) in Q16
770 w[i] = w[i] - nfa_qmul(lr, upd) - nfa_qmul(nfa_qmul(lr, wd), w[i])
771 i = i + 1
772 }
773 return 0
774}