nx_autograd_tensor.nx
buildroot/runtime/nx_autograd_tensor.nx
about
nx_autograd_tensor.nx -- TRAIN-R2 (T6): TENSOR reverse-mode autograd -- the identities that scale.
Rung-1's tape held scalars; real models hold ARRAYS. Same tape shape, now stride-7:
{op, ai, bi, rows, cols, valp, gradp} -- valp/gradp index a per-build BUMP ARENA of f32 cells (st[1]
counter; no allocator, no fragmentation, deterministic). Forward EAGER, backward ONE reverse sweep.
Four backward identities (learn these and every architecture MLP->transformer->DiT is composition):
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)
2. vadd y=x+b: both inputs receive gy unchanged (bias gradients are free)
3. relu elementwise: pass gy[i] iff the INPUT cell was > 0
4. mse L=(1/n)sum(p-t)^2: dp_i += gL*(2/n)*(p_i - t_i); target grad nobody reads
Scalar ops (sum, average) compose for free: matvec(1x1, 1x1) = scalar multiply, vadd of 1-cell nodes = add.
genealogy_id: linnainmaa_1970_reverse_mode_ad + the matmul-backward identity (dA=dC B^T, dB=A^T dC)
lineage_id: sovereign_tensor_tape_autograd_v1
license_tier: ORIGINAL verdict: GREEN (nx_train_r2_gate 2026-06-14: MLP gradcheck + affine-recover + bit-exact + AdamW)
dependencies 8 imports · 14 importers
diagram shows first 10 each side; +0 more imports, +4 more importers in the complete lists below.
imports: nx_f32.nxnx_f32_div.nxnx_f32_cvt.nxnx_f32_exp.nxnx_f32_log.nxnx_f32_sincos.nxnx_fft_f32.nxnx_syscalls.nx
imported by: nx_f32_attn_charlm_gate.nxnx_f32_miniqwen_train_gate.nxnx_f32_qwen2_train_gate.nxnx_fnet_mlm_gate.nxnx_fnet_model_gate.nxnx_fnet_train_gate.nxnx_reader_quantize_gate.nxnx_reader_squad_f32_gate.nxnx_ssm_gate.nxnx_ssm_lm_gate.nxnx_ta_parallel_grad_gate.nxnx_ta_transformer_gradcheck_gate.nxnx_train_r2_gate.nxnx_train_r3_gate.nx
call flow from main pre-order; caps 40 nodes / depth 6 declared; ↻ = already shown
structs
| none |
consts
| 23 | const TA_MAGIC_1000000: i64 = 1000000 |
| 24 | const TA_MAGIC_1024: i64 = 1024 |
| 26 | const TA_LEAF: i64 = 0 |
| 27 | const TA_MATVEC: i64 = 1 |
| 28 | const TA_VADD: i64 = 2 |
| 29 | const TA_RELU: i64 = 3 |
| 30 | const TA_MSE: i64 = 4 |
| 31 | const TA_SOFTMAX: i64 = 5 |
| 32 | const TA_SOFTCE: i64 = 6 |
| 33 | const TA_FNET: i64 = 7 |
| 34 | const TA_SSM: i64 = 8 |
| 35 | const TA_SLICE: i64 = 9 |
| 37 | const TA_MATMUL: i64 = 10 // C[m,p] = A[m,k] . B[k,p] |
| 38 | const TA_MATMUL_NT: i64 = 11 // S[m,p] = A[m,k] . B[p,k]^T (the Q.K^T contraction) |
| 39 | const TA_CMUL: i64 = 12 // y = a * c, c an f32 constant stored in bi (the 1/sqrt(d) scale) |
| 40 | const TA_SOFTMAX_ROWS: i64 = 13 // per-row softmax; bi=1 -> CAUSAL (row i over j<=i) |
| 41 | const TA_ROPE: i64 = 14 // rotary position embedding on [T,hd]; parameter-free, orthogonal |
| 42 | const TA_HADAMARD: i64 = 15 // y = a (*) b elementwise (the SwiGLU gate) |
| 43 | const TA_SILU: i64 = 16 // y = x*sigmoid(x) |
| 44 | const TA_RMSNORM_ROWS: i64 = 17 // per-row y = x/sqrt(mean(x^2)+eps) |
| 45 | const TA_SOFTCE_ROWS: i64 = 18 // fused per-row softmax-CE vs integer target ids (bi = ids ptr) |
| 46 | const TA_ROPE_TAB: i64 = 19 // RoPE from a PRECOMPUTED cos/sin table (bi = table ptr) -- nx_f32_sin/cos |
| 49 | const TA_F32_ZERO: i64 = 0 // IEEE-754 binary32 +0.0 |
| 50 | const TA_F32_ONE: i64 = 1065353216 // IEEE-754 binary32 1.0 |
functions
| 53 | func ta_constf(num: i64, den: i64) -> i64 { return nx_f32_div(nx_i32_to_f32(num), nx_i32_to_f32(den)) } called by 21: ta_det_initta_rmsnorm_rowsta_backwardta_f32_to_milliml_trainmo_train+15 calls 2: nx_f32_divnx_i32_to_f32 |
| 56 | func ta_det_init(dst: *i64, n: i64, seed: i64) -> i64 |
| 66 | func ta_new(tape: *i64, st: *i64, op: i64, ai: i64, bi: i64, rows: i64, cols: i64) -> i64 |
| 78 | func ta_leaf(tape: *i64, vals: *i64, st: *i64, rows: i64, cols: i64, src: *i64, soff: i64) -> i64 |
| 88 | func ta_matvec(tape: *i64, vals: *i64, st: *i64, aW: i64, bx: i64) -> i64 |
| 107 | func ta_vadd(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 |
| 121 | func ta_relu(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 |
| 139 | func ta_mse(tape: *i64, vals: *i64, st: *i64, pred: i64, target: i64) -> i64 |
| 157 | func ta_softmax(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 called by 2: g_smax_lossg_smax_grads calls 6: ta_newnx_f32_gtnx_f32_expnx_f32_subnx_f32_addnx_f32_div |
| 175 | func ta_softce(tape: *i64, vals: *i64, st: *i64, logits: i64, target: i64) -> i64 |
| 198 | func ta_fnet(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 |
| 219 | func ta_ssm(tape: *i64, vals: *i64, st: *i64, a_node: i64, x_node: i64) -> i64 |
| 243 | func ta_slice(tape: *i64, vals: *i64, st: *i64, x_node: i64, row: i64) -> i64 |
| 255 | func ta_matmul(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 |
| 278 | func ta_matmul_nt(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 |
| 301 | func ta_cmul(tape: *i64, vals: *i64, st: *i64, a: i64, c_f: i64) -> i64 |
| 313 | func ta_softmax_rows(tape: *i64, vals: *i64, st: *i64, a: i64, causal: i64) -> i64 |
| 339 | func ta_rope_ang(t: i64, i: i64, np: i64) -> i64 |
| 346 | func ta_rope(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 |
| 373 | func ta_rope_build_tab(tab: *i64, tmax: i64, np: i64) -> i64 |
| 390 | func ta_rope_tab(tape: *i64, vals: *i64, st: *i64, a: i64, tab: *i64) -> i64 |
| 415 | func ta_hadamard(tape: *i64, vals: *i64, st: *i64, a: i64, b: i64) -> i64 called by 8: rope_nodebuild_fwdrope_nodermsnorm_gammabuild_fwdqz_fwd_f32+2 calls 2: ta_newnx_f32_mul |
| 428 | func ta_sigmoid(x: i64) -> i64 |
| 433 | func ta_silu(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 called by 6: build_fwdbuild_fwdbuild_fwdqz_fwd_f32rf_fwdbld_silu calls 3: ta_newnx_f32_multa_sigmoid |
| 445 | func ta_rmsnorm_rows(tape: *i64, vals: *i64, st: *i64, a: i64) -> i64 |
| 467 | func ta_softce_rows(tape: *i64, vals: *i64, st: *i64, logits: i64, tgt: *i64) -> i64 called by 8: loss_ofmainloss_ofmainloss_ofmain+2 calls 8: ta_newnx_f32_gtnx_f32_addnx_f32_expnx_f32_subnx_f32_log+2 |
| 491 | func ta_val(tape: *i64, vals: *i64, k: i64, c: i64) -> i64 { return vals[tape[7 * k + 5] + c] } |
| 492 | func ta_grad(tape: *i64, grads: *i64, k: i64, c: i64) -> i64 { return grads[tape[7 * k + 6] + c] } |
| 495 | func ta_backward(tape: *i64, vals: *i64, grads: *i64, n: i64, root: i64) -> i64 called by 19: mainmainmainmainml_trainmo_train+13 calls 15: nx_f32_addnx_f32_mulnx_f32_gtnx_f32_divnx_i32_to_f32nx_f32_sub+9 |
| 899 | func ta_f32_to_milli(v: i64) -> i64 called by 16: gmmaingmmaingmmain+10 calls 6: nx_f32_mulnx_i32_to_f32nx_f32_ltnx_f32_negta_constfnx_f32_add |
| 917 | func main() -> i64 |