nx_autograd_tensor.nx source
↩ module page · 932 lines · 39566 B
1// nx_autograd_tensor.nx -- TRAIN-R2 (T6): TENSOR reverse-mode autograd -- the identities that scale.
2// Rung-1's tape held scalars; real models hold ARRAYS. Same tape shape, now stride-7:
3// {op, ai, bi, rows, cols, valp, gradp} -- valp/gradp index a per-build BUMP ARENA of f32 cells (st[1]
4// counter; no allocator, no fragmentation, deterministic). Forward EAGER, backward ONE reverse sweep.
5// Four backward identities (learn these and every architecture MLP->transformer->DiT is composition):
6// 1. matvec y=W*x (W r*c, x c): dW[i][j] += gy[i]*x[j] (outer product); dx[j] += sum_i W[i][j]*gy[i] (W^T gy)
7// 2. vadd y=x+b: both inputs receive gy unchanged (bias gradients are free)
8// 3. relu elementwise: pass gy[i] iff the INPUT cell was > 0
9// 4. mse L=(1/n)sum(p-t)^2: dp_i += gL*(2/n)*(p_i - t_i); target grad nobody reads
10// Scalar ops (sum, average) compose for free: matvec(1x1, 1x1) = scalar multiply, vadd of 1-cell nodes = add.
11//
12// genealogy_id: linnainmaa_1970_reverse_mode_ad + the matmul-backward identity (dA=dC B^T, dB=A^T dC)
13// lineage_id: sovereign_tensor_tape_autograd_v1
14// license_tier: ORIGINAL verdict: GREEN (nx_train_r2_gate 2026-06-14: MLP gradcheck + affine-recover + bit-exact + AdamW)
15import "nx_f32.nx" // add/sub/mul/neg/abs/lt/gt/sqrt -- eager forward + AdamW
16import "nx_f32_div.nx" // nx_f32_div
17import "nx_f32_cvt.nx" // nx_i32_to_f32
18import "nx_f32_exp.nx" // nx_f32_exp -- softmax / log-sum-exp
19import "nx_f32_log.nx" // nx_f32_log -- log-sum-exp
20import "nx_f32_sincos.nx" // nx_f32_sin/cos -- RoPE rotation (TA_ROPE)
21import "nx_fft_f32.nx" // fnet_mix_f32 -- the trainable FNet token mixer (TA_FNET op)
22import "nx_syscalls.nx" // sys_mmap
23const TA_MAGIC_1000000: i64 = 1000000
24const TA_MAGIC_1024: i64 = 1024
25
26const TA_LEAF: i64 = 0
27const TA_MATVEC: i64 = 1
28const TA_VADD: i64 = 2
29const TA_RELU: i64 = 3
30const TA_MSE: i64 = 4
31const TA_SOFTMAX: i64 = 5
32const TA_SOFTCE: i64 = 6
33const TA_FNET: i64 = 7
34const TA_SSM: i64 = 8
35const TA_SLICE: i64 = 9
36// ---- transformer ops (added 2026-07-09 for the R3e f32 neural reader; each mirrors its GREEN nfa_* twin) ----
37const TA_MATMUL: i64 = 10 // C[m,p] = A[m,k] . B[k,p]
38const TA_MATMUL_NT: i64 = 11 // S[m,p] = A[m,k] . B[p,k]^T (the Q.K^T contraction)
39const TA_CMUL: i64 = 12 // y = a * c, c an f32 constant stored in bi (the 1/sqrt(d) scale)
40const TA_SOFTMAX_ROWS: i64 = 13 // per-row softmax; bi=1 -> CAUSAL (row i over j<=i)
41const TA_ROPE: i64 = 14 // rotary position embedding on [T,hd]; parameter-free, orthogonal
42const TA_HADAMARD: i64 = 15 // y = a (*) b elementwise (the SwiGLU gate)
43const TA_SILU: i64 = 16 // y = x*sigmoid(x)
44const TA_RMSNORM_ROWS: i64 = 17 // per-row y = x/sqrt(mean(x^2)+eps)
45const TA_SOFTCE_ROWS: i64 = 18 // fused per-row softmax-CE vs integer target ids (bi = ids ptr)
46const TA_ROPE_TAB: i64 = 19 // RoPE from a PRECOMPUTED cos/sin table (bi = table ptr) -- nx_f32_sin/cos
47 // sys_mmap PER CALL, so per-element transcendentals in a training loop leak
48 // ~4KB/call -> OOM (measured: run-exit=137). Table = 2k calls ONCE.
49const TA_F32_ZERO: i64 = 0 // IEEE-754 binary32 +0.0
50const TA_F32_ONE: i64 = 1065353216 // IEEE-754 binary32 1.0
51
52// rational f32 constant num/den
53func ta_constf(num: i64, den: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(num), nx_i32_to_f32(den)) }
54
55// deterministic symmetry-breaking init: spread small f32 weights so relu units differ (zero-init = dead-relu trap).
56func ta_det_init(dst: *i64, n: i64, seed: i64) -> i64 {
57 var i: i64 = 0
58 while i < n {
59 dst[i] = ta_constf(((i * seed + 1) % 11) - 5, 8) // values in {-5/8 .. 5/8}
60 i = i + 1
61 }
62 return 0
63}
64
65// st[0] = next node index, st[1] = next free arena cell. Allocate a node, return its index.
66func ta_new(tape: *i64, st: *i64, op: i64, ai: i64, bi: i64, rows: i64, cols: i64) -> i64 {
67 let k: i64 = st[0]
68 let off: i64 = st[1]
69 tape[7 * k + 0] = op; tape[7 * k + 1] = ai; tape[7 * k + 2] = bi
70 tape[7 * k + 3] = rows; tape[7 * k + 4] = cols
71 tape[7 * k + 5] = off; tape[7 * k + 6] = off
72 st[0] = k + 1
73 st[1] = off + rows * cols
74 return k
75}
76
77// leaf holding rows*cols f32 cells copied from src[soff..].
78func ta_leaf(tape: *i64, vals: *i64, st: *i64, rows: i64, cols: i64, src: *i64, soff: i64) -> i64 {
79 let k: i64 = ta_new(tape, st, TA_LEAF, 0 - 1, 0 - 1, rows, cols)
80 let off: i64 = tape[7 * k + 5]
81 let sz: i64 = rows * cols
82 var i: i64 = 0
83 while i < sz { vals[off + i] = src[soff + i]; i = i + 1 }
84 return k
85}
86
87// y = W * x (W is r*c, x is c*1, y is r*1)
88func ta_matvec(tape: *i64, vals: *i64, st: *i64, aW: i64, bx: i64) -> i64 {
89 let r: i64 = tape[7 * aW + 3]
90 let c: i64 = tape[7 * aW + 4]
91 let k: i64 = ta_new(tape, st, TA_MATVEC, aW, bx, r, 1)
92 let offy: i64 = tape[7 * k + 5]
93 let offW: i64 = tape[7 * aW + 5]
94 let offx: i64 = tape[7 * bx + 5]
95 var i: i64 = 0
96 while i < r {
97 var acc: i64 = TA_F32_ZERO
98 var j: i64 = 0
99 while j < c { acc = nx_f32_add(acc, nx_f32_mul(vals[offW + i * c + j], vals[offx + j])); j = j + 1 }
100 vals[offy + i] = acc
101 i = i + 1
102 }
103 return k
104}
105
106// y = a + b (same shape)
107func ta_vadd(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
108 let r: i64 = tape[7 * a + 3]
109 let c: i64 = tape[7 * a + 4]
110 let k: i64 = ta_new(tape, st, TA_VADD, a, b, r, c)
111 let offy: i64 = tape[7 * k + 5]
112 let offa: i64 = tape[7 * a + 5]
113 let offb: i64 = tape[7 * b + 5]
114 let sz: i64 = r * c
115 var i: i64 = 0
116 while i < sz { vals[offy + i] = nx_f32_add(vals[offa + i], vals[offb + i]); i = i + 1 }
117 return k
118}
119
120// y = relu(a) elementwise
121func ta_relu(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
122 let r: i64 = tape[7 * a + 3]
123 let c: i64 = tape[7 * a + 4]
124 let k: i64 = ta_new(tape, st, TA_RELU, a, 0 - 1, r, c)
125 let offy: i64 = tape[7 * k + 5]
126 let offa: i64 = tape[7 * a + 5]
127 let sz: i64 = r * c
128 var i: i64 = 0
129 while i < sz {
130 var v: i64 = vals[offa + i]
131 if nx_f32_gt(v, TA_F32_ZERO) != 1 { v = TA_F32_ZERO }
132 vals[offy + i] = v
133 i = i + 1
134 }
135 return k
136}
137
138// L = (1/n) sum_i (pred_i - target_i)^2 (scalar, 1 cell)
139func ta_mse(tape: *i64, vals: *i64, st: *i64, pred: i64, target: i64) -> i64 {
140 let np: i64 = tape[7 * pred + 3] * tape[7 * pred + 4]
141 let k: i64 = ta_new(tape, st, TA_MSE, pred, target, 1, 1)
142 let offL: i64 = tape[7 * k + 5]
143 let offp: i64 = tape[7 * pred + 5]
144 let offt: i64 = tape[7 * target + 5]
145 var acc: i64 = TA_F32_ZERO
146 var i: i64 = 0
147 while i < np {
148 let di: i64 = nx_f32_sub(vals[offp + i], vals[offt + i])
149 acc = nx_f32_add(acc, nx_f32_mul(di, di))
150 i = i + 1
151 }
152 vals[offL] = nx_f32_mul(acc, nx_f32_div(nx_i32_to_f32(1), nx_i32_to_f32(np)))
153 return k
154}
155
156// y = softmax(x) over the node's cells (one vector). Stores y for the backward Jacobian-vector product.
157func ta_softmax(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
158 let n: i64 = tape[7 * a + 3] * tape[7 * a + 4]
159 let k: i64 = ta_new(tape, st, TA_SOFTMAX, a, 0 - 1, tape[7 * a + 3], tape[7 * a + 4])
160 let offy: i64 = tape[7 * k + 5]
161 let offa: i64 = tape[7 * a + 5]
162 var mx: i64 = vals[offa]
163 var i: i64 = 1
164 while i < n { if nx_f32_gt(vals[offa + i], mx) == 1 { mx = vals[offa + i] } i = i + 1 }
165 var sum: i64 = TA_F32_ZERO
166 i = 0
167 while i < n { let e: i64 = nx_f32_exp(nx_f32_sub(vals[offa + i], mx)); vals[offy + i] = e; sum = nx_f32_add(sum, e); i = i + 1 }
168 i = 0
169 while i < n { vals[offy + i] = nx_f32_div(vals[offy + i], sum); i = i + 1 }
170 return k
171}
172
173// L = softmax-cross-entropy(logits, target) = logsumexp(logits) - sum_i target_i*logits_i (scalar, stable via
174// max-subtract). The classifier loss head; backward identity is famously clean: dlogits_i = softmax_i - target_i.
175func ta_softce(tape: *i64, vals: *i64, st: *i64, logits: i64, target: i64) -> i64 {
176 let n: i64 = tape[7 * logits + 3] * tape[7 * logits + 4]
177 let k: i64 = ta_new(tape, st, TA_SOFTCE, logits, target, 1, 1)
178 let offL: i64 = tape[7 * k + 5]
179 let offp: i64 = tape[7 * logits + 5]
180 let offt: i64 = tape[7 * target + 5]
181 var mx: i64 = vals[offp]
182 var i: i64 = 1
183 while i < n { if nx_f32_gt(vals[offp + i], mx) == 1 { mx = vals[offp + i] } i = i + 1 }
184 var sume: i64 = TA_F32_ZERO
185 i = 0
186 while i < n { sume = nx_f32_add(sume, nx_f32_exp(nx_f32_sub(vals[offp + i], mx))); i = i + 1 }
187 let lse: i64 = nx_f32_add(mx, nx_f32_log(sume))
188 var dot: i64 = TA_F32_ZERO
189 i = 0
190 while i < n { dot = nx_f32_add(dot, nx_f32_mul(vals[offt + i], vals[offp + i])); i = i + 1 }
191 vals[offL] = nx_f32_sub(lse, dot)
192 return k
193}
194
195// FNet token-mixer as a tape op: y = Re(2D-DFT(x)); x and y are both [n,d] = [rows,cols]. The sub-quadratic
196// attention alternative, now differentiable. Forward = fnet_mix_f32; backward (self-adjoint) = the same mixer
197// applied to the upstream gradient (see ta_backward). ZERO parameters, zero stored activations.
198func ta_fnet(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
199 let n: i64 = tape[7 * a + 3]
200 let d: i64 = tape[7 * a + 4]
201 let k: i64 = ta_new(tape, st, TA_FNET, a, 0 - 1, n, d)
202 let offy: i64 = tape[7 * k + 5]
203 let offa: i64 = tape[7 * a + 5]
204 let sz: i64 = n * d
205 let re: *i64 = (sys_mmap(sz * 8)) as *i64
206 let im: *i64 = (sys_mmap(sz * 8)) as *i64
207 var i: i64 = 0
208 while i < sz { re[i] = vals[offa + i]; im[i] = TA_F32_ZERO; i = i + 1 }
209 fnet_mix_f32(re, im, n, d)
210 i = 0
211 while i < sz { vals[offy + i] = re[i]; i = i + 1 }
212 return k
213}
214
215// CAUSAL diagonal state-space mixer (the autoregressive sub-quadratic alternative to attention): per channel j,
216// h_t = a_j * h_{t-1} + x_{t,j}, y_{t,j} = h_t (h_{-1}=0).
217// y_t depends ONLY on x_0..x_t (causal) and costs O(n*d) (a linear scan). a_node = per-channel learnable decay
218// [d]; x_node = [n,d]; output [n,d]. Backward (in ta_backward) is backprop-through-time = a reverse scan.
219func ta_ssm(tape: *i64, vals: *i64, st: *i64, a_node: i64, x_node: i64) -> i64 {
220 let n: i64 = tape[7 * x_node + 3]
221 let d: i64 = tape[7 * x_node + 4]
222 let k: i64 = ta_new(tape, st, TA_SSM, a_node, x_node, n, d)
223 let offy: i64 = tape[7 * k + 5]
224 let offa: i64 = tape[7 * a_node + 5]
225 let offx: i64 = tape[7 * x_node + 5]
226 var j: i64 = 0
227 while j < d {
228 let aj: i64 = vals[offa + j]
229 var h: i64 = TA_F32_ZERO
230 var t: i64 = 0
231 while t < n {
232 h = nx_f32_add(nx_f32_mul(aj, h), vals[offx + t * d + j])
233 vals[offy + t * d + j] = h
234 t = t + 1
235 }
236 j = j + 1
237 }
238 return k
239}
240
241// extract row `row` of an [n,d] node as a [d] column vector. bi holds the integer row index (NOT a node).
242// Used to read each sequence position's state so a shared head can predict the next token (causal LM).
243func ta_slice(tape: *i64, vals: *i64, st: *i64, x_node: i64, row: i64) -> i64 {
244 let d: i64 = tape[7 * x_node + 4]
245 let k: i64 = ta_new(tape, st, TA_SLICE, x_node, row, d, 1)
246 let offy: i64 = tape[7 * k + 5]
247 let offx: i64 = tape[7 * x_node + 5]
248 var j: i64 = 0
249 while j < d { vals[offy + j] = vals[offx + row * d + j]; j = j + 1 }
250 return k
251}
252
253// ==== transformer-op forwards (R3e; f32 twins of the GREEN nfa_* ops) ====
254// C = A . B (A [m,k], B [k,p] -> C [m,p])
255func ta_matmul(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
256 let m: i64 = tape[7*a+3]
257 let kk: i64 = tape[7*a+4]
258 let p: i64 = tape[7*b+4]
259 let nd: i64 = ta_new(tape, st, TA_MATMUL, a, b, m, p)
260 let offC: i64 = tape[7*nd+5]
261 let offA: i64 = tape[7*a+5]
262 let offB: i64 = tape[7*b+5]
263 var i: i64 = 0
264 while i < m {
265 var j: i64 = 0
266 while j < p {
267 var acc: i64 = TA_F32_ZERO
268 var l: i64 = 0
269 while l < kk { acc = __f32_add(acc, __f32_mul(vals[offA + i*kk + l], vals[offB + l*p + j])); l = l + 1 } // HW-SSE (addss/mulss); bit-identical to soft-float, ~7x
270 vals[offC + i*p + j] = acc
271 j = j + 1
272 }
273 i = i + 1
274 }
275 return nd
276}
277// S = A . B^T (A [m,k], B [p,k] -> S [m,p])
278func ta_matmul_nt(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
279 let m: i64 = tape[7*a+3]
280 let kk: i64 = tape[7*a+4]
281 let p: i64 = tape[7*b+3]
282 let nd: i64 = ta_new(tape, st, TA_MATMUL_NT, a, b, m, p)
283 let offS: i64 = tape[7*nd+5]
284 let offA: i64 = tape[7*a+5]
285 let offB: i64 = tape[7*b+5]
286 var i: i64 = 0
287 while i < m {
288 var j: i64 = 0
289 while j < p {
290 var acc: i64 = TA_F32_ZERO
291 var l: i64 = 0
292 while l < kk { acc = __f32_add(acc, __f32_mul(vals[offA + i*kk + l], vals[offB + j*kk + l])); l = l + 1 } // HW-SSE; bit-identical, ~7x
293 vals[offS + i*p + j] = acc
294 j = j + 1
295 }
296 i = i + 1
297 }
298 return nd
299}
300// y = a * c (c = f32 constant in bi)
301func ta_cmul(tape: *i64, vals: *i64, st: *i64, a: i64, c_f: i64) -> i64 {
302 let r: i64 = tape[7*a+3]
303 let c: i64 = tape[7*a+4]
304 let n: i64 = r * c
305 let nd: i64 = ta_new(tape, st, TA_CMUL, a, c_f, r, c)
306 let offy: i64 = tape[7*nd+5]
307 let offa: i64 = tape[7*a+5]
308 var i: i64 = 0
309 while i < n { vals[offy+i] = nx_f32_mul(vals[offa+i], c_f); i = i + 1 }
310 return nd
311}
312// y = per-row softmax(a); causal=1 -> row i normalizes over j<=i only
313func ta_softmax_rows(tape: *i64, vals: *i64, st: *i64, a: i64, causal: i64) -> i64 {
314 let r: i64 = tape[7*a+3]
315 let c: i64 = tape[7*a+4]
316 let nd: i64 = ta_new(tape, st, TA_SOFTMAX_ROWS, a, causal, r, c)
317 let offy: i64 = tape[7*nd+5]
318 let offa: i64 = tape[7*a+5]
319 var i: i64 = 0
320 while i < r {
321 var lim: i64 = c
322 if causal == 1 { lim = i + 1 }
323 let base: i64 = i * c
324 var mx: i64 = vals[offa+base]
325 var j: i64 = 1
326 while j < lim { if nx_f32_gt(vals[offa+base+j], mx) == 1 { mx = vals[offa+base+j] } j = j + 1 }
327 var sum: i64 = TA_F32_ZERO
328 j = 0
329 while j < lim { let e: i64 = nx_f32_exp(nx_f32_sub(vals[offa+base+j], mx)); vals[offy+base+j] = e; sum = nx_f32_add(sum, e); j = j + 1 }
330 j = lim
331 while j < c { vals[offy+base+j] = TA_F32_ZERO; j = j + 1 }
332 j = 0
333 while j < lim { vals[offy+base+j] = nx_f32_div(vals[offy+base+j], sum); j = j + 1 }
334 i = i + 1
335 }
336 return nd
337}
338// RoPE angle for pair i of np at position t: t * base^(-i/np), base=1e6 (matches nfa/Qwen)
339func ta_rope_ang(t: i64, i: i64, np: i64) -> i64 {
340 let lnb: i64 = nx_f32_log(nx_i32_to_f32(TA_MAGIC_1000000))
341 let frac: i64 = nx_f32_div(nx_f32_mul(nx_i32_to_f32(i), lnb), nx_i32_to_f32(np))
342 let theta: i64 = nx_f32_exp(nx_f32_neg(frac))
343 return nx_f32_mul(nx_i32_to_f32(t), theta)
344}
345// y = RoPE(a) on [T,hd] (hd even): rotate each (2i,2i+1) pair by ang(t,i)
346func ta_rope(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
347 let T: i64 = tape[7*a+3]
348 let hd: i64 = tape[7*a+4]
349 let np: i64 = hd/2
350 let nd: i64 = ta_new(tape, st, TA_ROPE, a, 0 - 1, T, hd)
351 let offy: i64 = tape[7*nd+5]
352 let offa: i64 = tape[7*a+5]
353 var t: i64 = 0
354 while t < T {
355 var i: i64 = 0
356 while i < np {
357 let ang: i64 = ta_rope_ang(t, i, np)
358 let c: i64 = nx_f32_cos(ang)
359 let s: i64 = nx_f32_sin(ang)
360 let av: i64 = vals[offa + t*hd + 2*i]
361 let bv: i64 = vals[offa + t*hd + 2*i + 1]
362 vals[offy + t*hd + 2*i] = nx_f32_sub(nx_f32_mul(av,c), nx_f32_mul(bv,s))
363 vals[offy + t*hd + 2*i + 1] = nx_f32_add(nx_f32_mul(av,s), nx_f32_mul(bv,c))
364 i = i + 1
365 }
366 t = t + 1
367 }
368 return nd
369}
370// build the self-describing RoPE table: tab[0]=tmax, tab[1]=np, cos at tab[2 + t*np + i],
371// sin at tab[2 + tmax*np + t*np + i]. Calls nx_f32_sin/cos tmax*np times ONCE (they mmap per call -> never
372// call them inside a training loop). Caller mmaps tab with (2 + 2*tmax*np) cells.
373func ta_rope_build_tab(tab: *i64, tmax: i64, np: i64) -> i64 {
374 tab[0] = tmax
375 tab[1] = np
376 var t: i64 = 0
377 while t < tmax {
378 var i: i64 = 0
379 while i < np {
380 let ang: i64 = ta_rope_ang(t, i, np)
381 tab[2 + t*np + i] = nx_f32_cos(ang)
382 tab[2 + tmax*np + t*np + i] = nx_f32_sin(ang)
383 i = i + 1
384 }
385 t = t + 1
386 }
387 return 0
388}
389// y = RoPE(a) on [T,hd] from the precomputed table (T <= tab[0], hd/2 == tab[1])
390func ta_rope_tab(tape: *i64, vals: *i64, st: *i64, a: i64, tab: *i64) -> i64 {
391 let T: i64 = tape[7*a+3]
392 let hd: i64 = tape[7*a+4]
393 let np: i64 = hd/2
394 let tmax: i64 = tab[0]
395 let nd: i64 = ta_new(tape, st, TA_ROPE_TAB, a, tab as i64, T, hd)
396 let offy: i64 = tape[7*nd+5]
397 let offa: i64 = tape[7*a+5]
398 var t: i64 = 0
399 while t < T {
400 var i: i64 = 0
401 while i < np {
402 let c: i64 = tab[2 + t*np + i]
403 let s: i64 = tab[2 + tmax*np + t*np + i]
404 let av: i64 = vals[offa + t*hd + 2*i]
405 let bv: i64 = vals[offa + t*hd + 2*i + 1]
406 vals[offy + t*hd + 2*i] = nx_f32_sub(nx_f32_mul(av,c), nx_f32_mul(bv,s))
407 vals[offy + t*hd + 2*i + 1] = nx_f32_add(nx_f32_mul(av,s), nx_f32_mul(bv,c))
408 i = i + 1
409 }
410 t = t + 1
411 }
412 return nd
413}
414// y = a (*) b elementwise
415func ta_hadamard(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 {
416 let r: i64 = tape[7*a+3]
417 let c: i64 = tape[7*a+4]
418 let n: i64 = r * c
419 let nd: i64 = ta_new(tape, st, TA_HADAMARD, a, b, r, c)
420 let offy: i64 = tape[7*nd+5]
421 let offa: i64 = tape[7*a+5]
422 let offb: i64 = tape[7*b+5]
423 var i: i64 = 0
424 while i < n { vals[offy+i] = nx_f32_mul(vals[offa+i], vals[offb+i]); i = i + 1 }
425 return nd
426}
427// sigmoid(x) = 1/(1+exp(-x)) in f32
428func ta_sigmoid(x: i64) -> i64 {
429 let ex: i64 = nx_f32_exp(nx_f32_neg(x))
430 return nx_f32_div(TA_F32_ONE, nx_f32_add(TA_F32_ONE, ex))
431}
432// y = silu(a) = a*sigmoid(a)
433func ta_silu(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
434 let r: i64 = tape[7*a+3]
435 let c: i64 = tape[7*a+4]
436 let n: i64 = r * c
437 let nd: i64 = ta_new(tape, st, TA_SILU, a, 0 - 1, r, c)
438 let offy: i64 = tape[7*nd+5]
439 let offa: i64 = tape[7*a+5]
440 var i: i64 = 0
441 while i < n { let v: i64 = vals[offa+i]; vals[offy+i] = nx_f32_mul(v, ta_sigmoid(v)); i = i + 1 }
442 return nd
443}
444// y = rmsnorm of each row of [r,c]: y[i][j] = x[i][j]/sqrt(mean_j(x^2)+eps)
445func ta_rmsnorm_rows(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 {
446 let r: i64 = tape[7*a+3]
447 let c: i64 = tape[7*a+4]
448 let nd: i64 = ta_new(tape, st, TA_RMSNORM_ROWS, a, 0 - 1, r, c)
449 let offy: i64 = tape[7*nd+5]
450 let offa: i64 = tape[7*a+5]
451 let eps: i64 = ta_constf(1, TA_MAGIC_1000000)
452 var i: i64 = 0
453 while i < r {
454 let base: i64 = i * c
455 var ss: i64 = TA_F32_ZERO
456 var j: i64 = 0
457 while j < c { ss = nx_f32_add(ss, nx_f32_mul(vals[offa+base+j], vals[offa+base+j])); j = j + 1 }
458 let ms: i64 = nx_f32_div(ss, nx_i32_to_f32(c))
459 let sd: i64 = nx_f32_sqrt(nx_f32_add(ms, eps))
460 j = 0
461 while j < c { vals[offy+base+j] = nx_f32_div(vals[offa+base+j], sd); j = j + 1 }
462 i = i + 1
463 }
464 return nd
465}
466// L = (1/T) sum_t [ lse(logits[t]) - logits[t][id_t] ]; logits [T,V], tgt = *i64 of T ids (ptr in bi)
467func ta_softce_rows(tape: *i64, vals: *i64, st: *i64, logits: i64, tgt: *i64) -> i64 {
468 let T: i64 = tape[7*logits+3]
469 let V: i64 = tape[7*logits+4]
470 let nd: i64 = ta_new(tape, st, TA_SOFTCE_ROWS, logits, tgt as i64, 1, 1)
471 let offL: i64 = tape[7*nd+5]
472 let offp: i64 = tape[7*logits+5]
473 var total: i64 = TA_F32_ZERO
474 var t: i64 = 0
475 while t < T {
476 let base: i64 = t*V
477 var mx: i64 = vals[offp+base]
478 var j: i64 = 1
479 while j < V { if nx_f32_gt(vals[offp+base+j], mx) == 1 { mx = vals[offp+base+j] } j = j + 1 }
480 var sum: i64 = TA_F32_ZERO
481 j = 0
482 while j < V { sum = nx_f32_add(sum, nx_f32_exp(nx_f32_sub(vals[offp+base+j], mx))); j = j + 1 }
483 let lse: i64 = nx_f32_add(mx, nx_f32_log(sum))
484 total = nx_f32_add(total, nx_f32_sub(lse, vals[offp+base+tgt[t]]))
485 t = t + 1
486 }
487 vals[offL] = nx_f32_div(total, nx_i32_to_f32(T))
488 return nd
489}
490
491func ta_val(tape: *i64, vals: *i64, k: i64, c: i64) -> i64 { return vals[tape[7 * k + 5] + c] }
492func ta_grad(tape: *i64, grads: *i64, k: i64, c: i64) -> i64 { return grads[tape[7 * k + 6] + c] }
493
494// reverse sweep: zero all grads, seed grad[root]=1, accumulate the four identities backward.
495func ta_backward(tape: *i64, vals: *i64, grads: *i64, n: i64, root: i64) -> i64 {
496 var k: i64 = 0
497 while k < n {
498 let off: i64 = tape[7 * k + 5]
499 let sz: i64 = tape[7 * k + 3] * tape[7 * k + 4]
500 var i: i64 = 0
501 while i < sz { grads[off + i] = TA_F32_ZERO; i = i + 1 }
502 k = k + 1
503 }
504 grads[tape[7 * root + 6] + 0] = TA_F32_ONE
505 k = n - 1
506 while k >= 0 {
507 let op: i64 = tape[7 * k + 0]
508 if op == TA_MATVEC {
509 let aW: i64 = tape[7 * k + 1]
510 let bx: i64 = tape[7 * k + 2]
511 let r: i64 = tape[7 * k + 3]
512 let c: i64 = tape[7 * aW + 4]
513 let offgy: i64 = tape[7 * k + 6]
514 let offW: i64 = tape[7 * aW + 5]
515 let gW: i64 = tape[7 * aW + 6]
516 let offx: i64 = tape[7 * bx + 5]
517 let gx: i64 = tape[7 * bx + 6]
518 var i: i64 = 0
519 while i < r {
520 let gyi: i64 = grads[offgy + i]
521 var j: i64 = 0
522 while j < c {
523 grads[gW + i * c + j] = nx_f32_add(grads[gW + i * c + j], nx_f32_mul(gyi, vals[offx + j]))
524 grads[gx + j] = nx_f32_add(grads[gx + j], nx_f32_mul(vals[offW + i * c + j], gyi))
525 j = j + 1
526 }
527 i = i + 1
528 }
529 }
530 if op == TA_VADD {
531 let a: i64 = tape[7 * k + 1]
532 let b: i64 = tape[7 * k + 2]
533 let sz: i64 = tape[7 * k + 3] * tape[7 * k + 4]
534 let offgy: i64 = tape[7 * k + 6]
535 let ga: i64 = tape[7 * a + 6]
536 let gb: i64 = tape[7 * b + 6]
537 var i: i64 = 0
538 while i < sz {
539 let g: i64 = grads[offgy + i]
540 grads[ga + i] = nx_f32_add(grads[ga + i], g)
541 grads[gb + i] = nx_f32_add(grads[gb + i], g)
542 i = i + 1
543 }
544 }
545 if op == TA_RELU {
546 let a: i64 = tape[7 * k + 1]
547 let sz: i64 = tape[7 * k + 3] * tape[7 * k + 4]
548 let offgy: i64 = tape[7 * k + 6]
549 let ga: i64 = tape[7 * a + 6]
550 let offa: i64 = tape[7 * a + 5]
551 var i: i64 = 0
552 while i < sz {
553 if nx_f32_gt(vals[offa + i], TA_F32_ZERO) == 1 {
554 grads[ga + i] = nx_f32_add(grads[ga + i], grads[offgy + i])
555 }
556 i = i + 1
557 }
558 }
559 if op == TA_MSE {
560 let a: i64 = tape[7 * k + 1]
561 let b: i64 = tape[7 * k + 2]
562 let np: i64 = tape[7 * a + 3] * tape[7 * a + 4]
563 let offp: i64 = tape[7 * a + 5]
564 let gp: i64 = tape[7 * a + 6]
565 let offt: i64 = tape[7 * b + 5]
566 let gL: i64 = grads[tape[7 * k + 6] + 0]
567 let twon: i64 = nx_f32_div(nx_i32_to_f32(2), nx_i32_to_f32(np))
568 var i: i64 = 0
569 while i < np {
570 let di: i64 = nx_f32_sub(vals[offp + i], vals[offt + i])
571 grads[gp + i] = nx_f32_add(grads[gp + i], nx_f32_mul(gL, nx_f32_mul(twon, di)))
572 i = i + 1
573 }
574 }
575 if op == TA_SOFTMAX {
576 let a: i64 = tape[7 * k + 1]
577 let n: i64 = tape[7 * k + 3] * tape[7 * k + 4]
578 let offy: i64 = tape[7 * k + 5]
579 let offgy: i64 = tape[7 * k + 6]
580 let ga: i64 = tape[7 * a + 6]
581 var dot: i64 = TA_F32_ZERO
582 var j: i64 = 0
583 while j < n { dot = nx_f32_add(dot, nx_f32_mul(vals[offy + j], grads[offgy + j])); j = j + 1 }
584 var i: i64 = 0
585 while i < n {
586 grads[ga + i] = nx_f32_add(grads[ga + i], nx_f32_mul(vals[offy + i], nx_f32_sub(grads[offgy + i], dot)))
587 i = i + 1
588 }
589 }
590 if op == TA_SOFTCE {
591 let a: i64 = tape[7 * k + 1]
592 let b: i64 = tape[7 * k + 2]
593 let n: i64 = tape[7 * a + 3] * tape[7 * a + 4]
594 let offp: i64 = tape[7 * a + 5]
595 let gp: i64 = tape[7 * a + 6]
596 let offt: i64 = tape[7 * b + 5]
597 let gL: i64 = grads[tape[7 * k + 6] + 0]
598 var mx: i64 = vals[offp]
599 var i: i64 = 1
600 while i < n { if nx_f32_gt(vals[offp + i], mx) == 1 { mx = vals[offp + i] } i = i + 1 }
601 var sum: i64 = TA_F32_ZERO
602 i = 0
603 while i < n { sum = nx_f32_add(sum, nx_f32_exp(nx_f32_sub(vals[offp + i], mx))); i = i + 1 }
604 i = 0
605 while i < n {
606 let smi: i64 = nx_f32_div(nx_f32_exp(nx_f32_sub(vals[offp + i], mx)), sum)
607 grads[gp + i] = nx_f32_add(grads[gp + i], nx_f32_mul(gL, nx_f32_sub(smi, vals[offt + i])))
608 i = i + 1
609 }
610 }
611 if op == TA_FNET {
612 let a: i64 = tape[7 * k + 1]
613 let n: i64 = tape[7 * k + 3]
614 let d: i64 = tape[7 * k + 4]
615 let sz: i64 = n * d
616 let offgy: i64 = tape[7 * k + 6]
617 let ga: i64 = tape[7 * a + 6]
618 let re: *i64 = (sys_mmap(sz * 8)) as *i64
619 let im: *i64 = (sys_mmap(sz * 8)) as *i64
620 var i: i64 = 0
621 while i < sz { re[i] = grads[offgy + i]; im[i] = TA_F32_ZERO; i = i + 1 }
622 fnet_mix_f32(re, im, n, d)
623 i = 0
624 while i < sz { grads[ga + i] = nx_f32_add(grads[ga + i], re[i]); i = i + 1 }
625 }
626 if op == TA_SSM {
627 let an: i64 = tape[7 * k + 1]
628 let xn: i64 = tape[7 * k + 2]
629 let n: i64 = tape[7 * k + 3]
630 let d: i64 = tape[7 * k + 4]
631 let offy: i64 = tape[7 * k + 5]
632 let offgy: i64 = tape[7 * k + 6]
633 let offa: i64 = tape[7 * an + 5]
634 let ga: i64 = tape[7 * an + 6]
635 let gx: i64 = tape[7 * xn + 6]
636 var j: i64 = 0
637 while j < d {
638 let aj: i64 = vals[offa + j]
639 var s: i64 = TA_F32_ZERO
640 var t: i64 = n - 1
641 while t >= 0 {
642 s = nx_f32_add(grads[offgy + t * d + j], nx_f32_mul(aj, s)) // s_t = g_t + a*s_{t+1}
643 grads[gx + t * d + j] = nx_f32_add(grads[gx + t * d + j], s) // dx_t = s_t
644 var hprev: i64 = TA_F32_ZERO
645 if t > 0 { hprev = vals[offy + (t - 1) * d + j] }
646 grads[ga + j] = nx_f32_add(grads[ga + j], nx_f32_mul(s, hprev)) // da += s_t * h_{t-1}
647 t = t - 1
648 }
649 j = j + 1
650 }
651 }
652 if op == TA_SLICE {
653 let xn: i64 = tape[7 * k + 1]
654 let row: i64 = tape[7 * k + 2]
655 let d: i64 = tape[7 * k + 3]
656 let offgy: i64 = tape[7 * k + 6]
657 let gx: i64 = tape[7 * xn + 6]
658 var j: i64 = 0
659 while j < d { grads[gx + row * d + j] = nx_f32_add(grads[gx + row * d + j], grads[offgy + j]); j = j + 1 }
660 }
661 // ==== transformer-op backwards (R3e; identities verbatim from the GREEN nfa_* lib, f32 arithmetic) ====
662 if op == TA_MATMUL {
663 // 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]
664 let a: i64 = tape[7*k+1]
665 let b: i64 = tape[7*k+2]
666 let m: i64 = tape[7*k+3]
667 let p: i64 = tape[7*k+4]
668 let kk: i64 = tape[7*a+4]
669 let offgy: i64 = tape[7*k+6]
670 let offA: i64 = tape[7*a+5]
671 let gA: i64 = tape[7*a+6]
672 let offB: i64 = tape[7*b+5]
673 let gB: i64 = tape[7*b+6]
674 var i: i64 = 0
675 while i < m {
676 var j: i64 = 0
677 while j < p {
678 let g: i64 = grads[offgy + i*p + j]
679 var l: i64 = 0
680 while l < kk {
681 grads[gA + i*kk + l] = __f32_add(grads[gA + i*kk + l], __f32_mul(g, vals[offB + l*p + j]))
682 grads[gB + l*p + j] = __f32_add(grads[gB + l*p + j], __f32_mul(vals[offA + i*kk + l], g))
683 l = l + 1
684 }
685 j = j + 1
686 }
687 i = i + 1
688 }
689 }
690 if op == TA_MATMUL_NT {
691 // 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]
692 let a: i64 = tape[7*k+1]
693 let b: i64 = tape[7*k+2]
694 let m: i64 = tape[7*k+3]
695 let p: i64 = tape[7*k+4]
696 let kk: i64 = tape[7*a+4]
697 let offgy: i64 = tape[7*k+6]
698 let offA: i64 = tape[7*a+5]
699 let gA: i64 = tape[7*a+6]
700 let offB: i64 = tape[7*b+5]
701 let gB: i64 = tape[7*b+6]
702 var i: i64 = 0
703 while i < m {
704 var j: i64 = 0
705 while j < p {
706 let g: i64 = grads[offgy + i*p + j]
707 var l: i64 = 0
708 while l < kk {
709 grads[gA + i*kk + l] = __f32_add(grads[gA + i*kk + l], __f32_mul(g, vals[offB + j*kk + l]))
710 grads[gB + j*kk + l] = __f32_add(grads[gB + j*kk + l], __f32_mul(g, vals[offA + i*kk + l]))
711 l = l + 1
712 }
713 j = j + 1
714 }
715 i = i + 1
716 }
717 }
718 if op == TA_CMUL {
719 let a: i64 = tape[7*k+1]
720 let c_f: i64 = tape[7*k+2]
721 let n2: i64 = tape[7*k+3] * tape[7*k+4]
722 let offgy: i64 = tape[7*k+6]
723 let ga: i64 = tape[7*a+6]
724 var i: i64 = 0
725 while i < n2 { grads[ga+i] = nx_f32_add(grads[ga+i], nx_f32_mul(grads[offgy+i], c_f)); i = i + 1 }
726 }
727 if op == TA_SOFTMAX_ROWS {
728 // per-row JVP over j<lim (causal: lim=i+1); masked entries -> 0 grad
729 let a: i64 = tape[7*k+1]
730 let causal: i64 = tape[7*k+2]
731 let r: i64 = tape[7*k+3]
732 let c: i64 = tape[7*k+4]
733 let offy: i64 = tape[7*k+5]
734 let offgy: i64 = tape[7*k+6]
735 let ga: i64 = tape[7*a+6]
736 var i: i64 = 0
737 while i < r {
738 var lim: i64 = c
739 if causal == 1 { lim = i + 1 }
740 let base: i64 = i * c
741 var dot: i64 = TA_F32_ZERO
742 var j: i64 = 0
743 while j < lim { dot = nx_f32_add(dot, nx_f32_mul(vals[offy+base+j], grads[offgy+base+j])); j = j + 1 }
744 j = 0
745 while j < lim { grads[ga+base+j] = nx_f32_add(grads[ga+base+j], nx_f32_mul(vals[offy+base+j], nx_f32_sub(grads[offgy+base+j], dot))); j = j + 1 }
746 i = i + 1
747 }
748 }
749 if op == TA_ROPE {
750 // backward = rotate gy by -ang (orthogonal): gv0=g0*c+g1*s ; gv1=g1*c-g0*s
751 let a: i64 = tape[7*k+1]
752 let T: i64 = tape[7*k+3]
753 let hd: i64 = tape[7*k+4]
754 let np: i64 = hd/2
755 let offgy: i64 = tape[7*k+6]
756 let ga: i64 = tape[7*a+6]
757 var t: i64 = 0
758 while t < T {
759 var i: i64 = 0
760 while i < np {
761 let ang: i64 = ta_rope_ang(t, i, np)
762 let c: i64 = nx_f32_cos(ang)
763 let s: i64 = nx_f32_sin(ang)
764 let g0: i64 = grads[offgy + t*hd + 2*i]
765 let g1: i64 = grads[offgy + t*hd + 2*i + 1]
766 grads[ga + t*hd + 2*i] = nx_f32_add(grads[ga + t*hd + 2*i], nx_f32_add(nx_f32_mul(g0,c), nx_f32_mul(g1,s)))
767 grads[ga + t*hd + 2*i + 1] = nx_f32_add(grads[ga + t*hd + 2*i + 1], nx_f32_sub(nx_f32_mul(g1,c), nx_f32_mul(g0,s)))
768 i = i + 1
769 }
770 t = t + 1
771 }
772 }
773 if op == TA_HADAMARD {
774 let a: i64 = tape[7*k+1]
775 let b: i64 = tape[7*k+2]
776 let n2: i64 = tape[7*k+3] * tape[7*k+4]
777 let offgy: i64 = tape[7*k+6]
778 let offa: i64 = tape[7*a+5]
779 let offb: i64 = tape[7*b+5]
780 let ga: i64 = tape[7*a+6]
781 let gb: i64 = tape[7*b+6]
782 var i: i64 = 0
783 while i < n2 {
784 let g: i64 = grads[offgy+i]
785 grads[ga+i] = nx_f32_add(grads[ga+i], nx_f32_mul(g, vals[offb+i]))
786 grads[gb+i] = nx_f32_add(grads[gb+i], nx_f32_mul(g, vals[offa+i]))
787 i = i + 1
788 }
789 }
790 if op == TA_ROPE_TAB {
791 // backward = rotate gy by -ang from the table: gv0=g0*c+g1*s ; gv1=g1*c-g0*s
792 let a: i64 = tape[7*k+1]
793 let tab: *i64 = tape[7*k+2] as *i64
794 let T: i64 = tape[7*k+3]
795 let hd: i64 = tape[7*k+4]
796 let np: i64 = hd/2
797 let tmax: i64 = tab[0]
798 let offgy: i64 = tape[7*k+6]
799 let ga: i64 = tape[7*a+6]
800 var t: i64 = 0
801 while t < T {
802 var i: i64 = 0
803 while i < np {
804 let c: i64 = tab[2 + t*np + i]
805 let s: i64 = tab[2 + tmax*np + t*np + i]
806 let g0: i64 = grads[offgy + t*hd + 2*i]
807 let g1: i64 = grads[offgy + t*hd + 2*i + 1]
808 grads[ga + t*hd + 2*i] = nx_f32_add(grads[ga + t*hd + 2*i], nx_f32_add(nx_f32_mul(g0,c), nx_f32_mul(g1,s)))
809 grads[ga + t*hd + 2*i + 1] = nx_f32_add(grads[ga + t*hd + 2*i + 1], nx_f32_sub(nx_f32_mul(g1,c), nx_f32_mul(g0,s)))
810 i = i + 1
811 }
812 t = t + 1
813 }
814 }
815 if op == TA_SILU {
816 // silu'(x) = s + x*s*(1-s), s = sigmoid(x)
817 let a: i64 = tape[7*k+1]
818 let n2: i64 = tape[7*k+3] * tape[7*k+4]
819 let offgy: i64 = tape[7*k+6]
820 let offa: i64 = tape[7*a+5]
821 let ga: i64 = tape[7*a+6]
822 var i: i64 = 0
823 while i < n2 {
824 let x: i64 = vals[offa+i]
825 let s: i64 = ta_sigmoid(x)
826 let d: i64 = nx_f32_add(s, nx_f32_mul(x, nx_f32_mul(s, nx_f32_sub(TA_F32_ONE, s))))
827 grads[ga+i] = nx_f32_add(grads[ga+i], nx_f32_mul(grads[offgy+i], d))
828 i = i + 1
829 }
830 }
831 if op == TA_RMSNORM_ROWS {
832 // per row: dx_j = gy_j/sd - y_j*dot/(c*(ms+eps)), dot = sum_j gy_j x_j (recompute fwd stats)
833 let a: i64 = tape[7*k+1]
834 let r: i64 = tape[7*k+3]
835 let c: i64 = tape[7*k+4]
836 let offy: i64 = tape[7*k+5]
837 let offgy: i64 = tape[7*k+6]
838 let offa: i64 = tape[7*a+5]
839 let ga: i64 = tape[7*a+6]
840 let eps: i64 = ta_constf(1, TA_MAGIC_1000000)
841 let cf: i64 = nx_i32_to_f32(c)
842 var i: i64 = 0
843 while i < r {
844 let base: i64 = i * c
845 var ss: i64 = TA_F32_ZERO
846 var j: i64 = 0
847 while j < c { ss = nx_f32_add(ss, nx_f32_mul(vals[offa+base+j], vals[offa+base+j])); j = j + 1 }
848 let ms: i64 = nx_f32_add(nx_f32_div(ss, cf), eps)
849 let sd: i64 = nx_f32_sqrt(ms)
850 var dot: i64 = TA_F32_ZERO
851 j = 0
852 while j < c { dot = nx_f32_add(dot, nx_f32_mul(grads[offgy+base+j], vals[offa+base+j])); j = j + 1 }
853 let denom: i64 = nx_f32_mul(cf, ms)
854 j = 0
855 while j < c {
856 let t1: i64 = nx_f32_div(grads[offgy+base+j], sd)
857 let t2: i64 = nx_f32_div(nx_f32_mul(vals[offy+base+j], dot), denom)
858 grads[ga+base+j] = nx_f32_add(grads[ga+base+j], nx_f32_sub(t1, t2))
859 j = j + 1
860 }
861 i = i + 1
862 }
863 }
864 if op == TA_SOFTCE_ROWS {
865 // dlogit[t][j] = gL * (softmax[t][j] - onehot[t][j]) / T (exact fused identity)
866 let logits: i64 = tape[7*k+1]
867 let tgt: *i64 = tape[7*k+2] as *i64
868 let T: i64 = tape[7*logits+3]
869 let V: i64 = tape[7*logits+4]
870 let offp: i64 = tape[7*logits+5]
871 let gp: i64 = tape[7*logits+6]
872 let gL: i64 = grads[tape[7*k+6] + 0]
873 let invT: i64 = nx_f32_div(TA_F32_ONE, nx_i32_to_f32(T))
874 var t: i64 = 0
875 while t < T {
876 let base: i64 = t*V
877 var mx: i64 = vals[offp+base]
878 var j: i64 = 1
879 while j < V { if nx_f32_gt(vals[offp+base+j], mx) == 1 { mx = vals[offp+base+j] } j = j + 1 }
880 var sum: i64 = TA_F32_ZERO
881 j = 0
882 while j < V { sum = nx_f32_add(sum, nx_f32_exp(nx_f32_sub(vals[offp+base+j], mx))); j = j + 1 }
883 j = 0
884 while j < V {
885 var smi: i64 = nx_f32_div(nx_f32_exp(nx_f32_sub(vals[offp+base+j], mx)), sum)
886 if j == tgt[t] { smi = nx_f32_sub(smi, TA_F32_ONE) }
887 grads[gp+base+j] = nx_f32_add(grads[gp+base+j], nx_f32_mul(gL, nx_f32_mul(smi, invT)))
888 j = j + 1
889 }
890 t = t + 1
891 }
892 }
893 k = k - 1
894 }
895 return 0
896}
897
898// f32 -> round(v*1000) signed int, for readable logging (nearest by linear advance over the i32->f32 ladder).
899func ta_f32_to_milli(v: i64) -> i64 {
900 var s: i64 = nx_f32_mul(v, nx_i32_to_f32(1000))
901 var neg: i64 = 0
902 if nx_f32_lt(s, TA_F32_ZERO) == 1 { neg = 1; s = nx_f32_neg(s) }
903 let half: i64 = ta_constf(1, 2)
904 var m: i64 = 0
905 var go: i64 = 1
906 while go == 1 {
907 if nx_f32_lt(nx_f32_add(nx_i32_to_f32(m), half), s) == 1 {
908 m = m + 1
909 if m >= TA_MAGIC_1000000 { go = 0 }
910 } else { go = 0 }
911 }
912 if neg == 1 { return 0 - m }
913 return m
914}
915
916// compile-only smoke: y = W*x, W=[[1,2],[3,4]], x=[1,1] -> y=[3,7]; backward runs.
917func main() -> i64 {
918 let tape: *i64 = (sys_mmap(256 * 7 * 8)) as *i64
919 let vals: *i64 = (sys_mmap(TA_MAGIC_1024 * 8)) as *i64
920 let grads: *i64 = (sys_mmap(TA_MAGIC_1024 * 8)) as *i64
921 let st: *i64 = (sys_mmap(2 * 8)) as *i64
922 st[0] = 0; st[1] = 0
923 let wsrc: *i64 = (sys_mmap(4 * 8)) as *i64
924 wsrc[0] = nx_i32_to_f32(1); wsrc[1] = nx_i32_to_f32(2); wsrc[2] = nx_i32_to_f32(3); wsrc[3] = nx_i32_to_f32(4)
925 let xsrc: *i64 = (sys_mmap(2 * 8)) as *i64
926 xsrc[0] = nx_i32_to_f32(1); xsrc[1] = nx_i32_to_f32(1)
927 let W: i64 = ta_leaf(tape, vals, st, 2, 2, wsrc, 0)
928 let x: i64 = ta_leaf(tape, vals, st, 2, 1, xsrc, 0)
929 let y: i64 = ta_matvec(tape, vals, st, W, x)
930 ta_backward(tape, vals, grads, st[0], y)
931 return 0
932}