`nx_ssm_gate.nx` source

↩ module page · 177 lines · 7799 B
1// nx_ssm_gate.nx -- GATE for SSM-001: the CAUSAL state-space mixer (the autoregressive sub-quadratic lane,
2// vs FNet's bidirectional encoder lane). Proves, by RUNNING:
3//   G1  GRADCHECK through the autograd: x[4,2] + learnable decay a[2] -> ta_ssm -> mse(.,target); analytic
4//       gradients for BOTH the input dx AND the decay da (backprop-through-time) vs central finite difference
5//       (h=1/128, rel<1/32 floor 1/64). Proves the reverse-scan backward is correct.
6//   G2  CAUSALITY: perturbing the input at the LAST time step leaves every EARLIER output unchanged, and
7//       changes only the last output -- i.e. the operator's Jacobian is lower-triangular (y_t sees only x_<=t).
8//       This is what FNet cannot do and what autoregressive generation requires.
9//
10// Evidence -> knowledge/status/ssm.log (SSMGATE authored=organ ... verdict=GREEN). license_tier: ORIGINAL
11import "nx_autograd_tensor.nx"
12import "nx_syscalls.nx"
13
14const SS_N: i64 = 4
15const SS_D: i64 = 2
16const SS_ND: i64 = 8
17const SS_LOG: *u8 = "knowledge/status/ssm.log"
18
19func ss_w(fd: i64, s: *u8) -> i64 { var n: i64 = 0; while s[n] != (0 as u8) { n = n + 1 } sys_write(fd, s, n); return 0 }
20func ss_wn(fd: i64, v: i64) -> i64 {
  let bb: *u8 = sys_mmap(28); var m: i64 = v
  if m < 0 { m = 0 - m; sys_write(fd, "-" as *u8, 1) }
  let t: *u8 = sys_mmap(28); var k: i64 = 0
  if m == 0 { t[0] = 48; k = 1 }
  while m > 0 { t[k] = (48 + (m % 10)) as u8; m = m / 10; k = k + 1 }
  var i: i64 = 0
  while i < k { bb[i] = t[k - 1 - i]; i = i + 1 }
  sys_write(fd, bb, k); return 0
29}
30
31func ss_build(tape: *i64, vals: *i64, st: *i64, a: *i64, x: *i64, target: *i64, ax: *i64) -> i64 {
  st[0] = 0; st[1] = 0
  let an: i64 = ta_leaf(tape, vals, st, SS_D, 1, a, 0)
  let xn: i64 = ta_leaf(tape, vals, st, SS_N, SS_D, x, 0)
  let y:  i64 = ta_ssm(tape, vals, st, an, xn)
  let tn: i64 = ta_leaf(tape, vals, st, SS_N, SS_D, target, 0)
  let loss: i64 = ta_mse(tape, vals, st, y, tn)
  ax[0] = an; ax[1] = xn
  return loss
40}
41func ss_loss(tape: *i64, vals: *i64, st: *i64, a: *i64, x: *i64, target: *i64) -> i64 {
  let ax: *i64 = (sys_mmap(2 * 8)) as *i64
  let loss: i64 = ss_build(tape, vals, st, a, x, target, ax)
  return ta_val(tape, vals, loss, 0)
45}
46// forward only: write all SS_ND output cells of ta_ssm(a,x) into yout
47func ss_forward(tape: *i64, vals: *i64, st: *i64, a: *i64, x: *i64, yout: *i64) -> i64 {
  st[0] = 0; st[1] = 0
  let an: i64 = ta_leaf(tape, vals, st, SS_D, 1, a, 0)
  let xn: i64 = ta_leaf(tape, vals, st, SS_N, SS_D, x, 0)
  let y:  i64 = ta_ssm(tape, vals, st, an, xn)
  var i: i64 = 0
  while i < SS_ND { yout[i] = ta_val(tape, vals, y, i); i = i + 1 }
  return 0
55}
56
57func main() -> i64 {
  var ok: i64 = 1
  let tape: *i64 = (sys_mmap(256 * 7 * 8)) as *i64
  let vals: *i64 = (sys_mmap(4096 * 8)) as *i64
  let grads: *i64 = (sys_mmap(4096 * 8)) as *i64
  let st: *i64 = (sys_mmap(2 * 8)) as *i64
63
  let a: *i64 = (sys_mmap(SS_D * 8)) as *i64
  a[0] = ta_constf(1, 2); a[1] = ta_constf(7, 10)        // decays 0.5, 0.7 (stable)
  let x: *i64 = (sys_mmap(SS_ND * 8)) as *i64
  let target: *i64 = (sys_mmap(SS_ND * 8)) as *i64
  var i: i64 = 0
  while i < SS_ND {
      x[i] = ta_constf((i % 5) - 2, 2)        // -1 .. 1
      target[i] = ta_constf((i % 7) - 3, 4)    // -0.75 .. 0.75
      i = i + 1
  }
74
  // ---------- G1: gradcheck dx and da ----------
  let ax: *i64 = (sys_mmap(2 * 8)) as *i64
  let loss: i64 = ss_build(tape, vals, st, a, x, target, ax)
  ta_backward(tape, vals, grads, st[0], loss)
  let ana_dx: *i64 = (sys_mmap(SS_ND * 8)) as *i64
  let ana_da: *i64 = (sys_mmap(SS_D * 8)) as *i64
  i = 0
  while i < SS_ND { ana_dx[i] = ta_grad(tape, grads, ax[1], i); i = i + 1 }
  i = 0
  while i < SS_D { ana_da[i] = ta_grad(tape, grads, ax[0], i); i = i + 1 }
85
  let h: i64 = ta_constf(1, 128)
  let flo: i64 = ta_constf(1, 64)
  let tol: i64 = ta_constf(1, 32)
  var gcPass: i64 = 1
  var gcWorst: i64 = 0
  // dx gradcheck
  let xp: *i64 = (sys_mmap(SS_ND * 8)) as *i64
  let xm: *i64 = (sys_mmap(SS_ND * 8)) as *i64
  var c: i64 = 0
  while c < SS_ND {
      var j: i64 = 0
      while j < SS_ND { xp[j] = x[j]; xm[j] = x[j]; j = j + 1 }
      xp[c] = nx_f32_add(x[c], h); xm[c] = nx_f32_sub(x[c], h)
      let lp: i64 = ss_loss(tape, vals, st, a, xp, target)
      let lm: i64 = ss_loss(tape, vals, st, a, xm, target)
      let fd: i64 = nx_f32_div(nx_f32_sub(lp, lm), nx_f32_add(h, h))
      let num: i64 = nx_f32_abs(nx_f32_sub(fd, ana_dx[c]))
      var den: i64 = nx_f32_abs(ana_dx[c])
      if nx_f32_lt(den, flo) == 1 { den = flo }
      if nx_f32_lt(num, nx_f32_mul(tol, den)) != 1 { gcPass = 0 }
      let nm: i64 = ta_f32_to_milli(num)
      if nm > gcWorst { gcWorst = nm }
      c = c + 1
  }
  // da gradcheck
  let ap: *i64 = (sys_mmap(SS_D * 8)) as *i64
  let am: *i64 = (sys_mmap(SS_D * 8)) as *i64
  c = 0
  while c < SS_D {
      var j: i64 = 0
      while j < SS_D { ap[j] = a[j]; am[j] = a[j]; j = j + 1 }
      ap[c] = nx_f32_add(a[c], h); am[c] = nx_f32_sub(a[c], h)
      let lp: i64 = ss_loss(tape, vals, st, ap, x, target)
      let lm: i64 = ss_loss(tape, vals, st, am, x, target)
      let fd: i64 = nx_f32_div(nx_f32_sub(lp, lm), nx_f32_add(h, h))
      let num: i64 = nx_f32_abs(nx_f32_sub(fd, ana_da[c]))
      var den: i64 = nx_f32_abs(ana_da[c])
      if nx_f32_lt(den, flo) == 1 { den = flo }
      if nx_f32_lt(num, nx_f32_mul(tol, den)) != 1 { gcPass = 0 }
      let nm: i64 = ta_f32_to_milli(num)
      if nm > gcWorst { gcWorst = nm }
      c = c + 1
  }
  if gcPass != 1 { ok = 0 }
130
  // ---------- G2: causality (perturb last time step; earlier outputs must not change) ----------
  let y1: *i64 = (sys_mmap(SS_ND * 8)) as *i64
  let y2: *i64 = (sys_mmap(SS_ND * 8)) as *i64
  ss_forward(tape, vals, st, a, x, y1)
  let x2: *i64 = (sys_mmap(SS_ND * 8)) as *i64
  i = 0
  while i < SS_ND { x2[i] = x[i]; i = i + 1 }
  var jj: i64 = 0
  while jj < SS_D { x2[(SS_N - 1) * SS_D + jj] = nx_f32_add(x2[(SS_N - 1) * SS_D + jj], ta_constf(1, 2)); jj = jj + 1 }
  ss_forward(tape, vals, st, a, x2, y2)
  var causalPass: i64 = 1
  var earlier_changed: i64 = 0
  var last_changed: i64 = 0
  var t: i64 = 0
  while t < SS_N {
      var j: i64 = 0
      while j < SS_D {
          let same: i64 = nx_f32_eq(y1[t * SS_D + j], y2[t * SS_D + j])
          if t < SS_N - 1 { if same != 1 { earlier_changed = 1 } }
          else { if same != 1 { last_changed = 1 } }
          j = j + 1
      }
      t = t + 1
  }
  if earlier_changed != 0 { causalPass = 0 }   // future input must NOT affect past outputs
  if last_changed != 1 { causalPass = 0 }       // but it MUST affect the current output
  if causalPass != 1 { ok = 0 }
158
  var fdi: i64 = 1
  while fdi >= 0 {
      var out: i64 = 1
      if fdi == 0 { out = sys_openat_append(SS_LOG, 420) }
      if out >= 0 {
          ss_w(out, "SSMGATE authored=organ op=causal-state-space-scan mixer=autoregressive-subquadratic" as *u8)
          ss_w(out, " | G1_gradcheck_dx+da_pass=" as *u8); ss_wn(out, gcPass)
          ss_w(out, " worst_|fd-analytic|_milli=" as *u8); ss_wn(out, gcWorst)
          ss_w(out, " | G2_causality_pass=" as *u8); ss_wn(out, causalPass)
          ss_w(out, " earlier_changed=" as *u8); ss_wn(out, earlier_changed)
          ss_w(out, " last_changed=" as *u8); ss_wn(out, last_changed)
          if ok == 1 { ss_w(out, " verdict=GREEN\n" as *u8) } else { ss_w(out, " verdict=RED\n" as *u8) }
          if fdi == 0 { sys_close(out) }
      }
      fdi = fdi - 1
  }
  if ok == 1 { return 0 }
  return 1
177}
nx_ssm_gate.nx source

`nx_ssm_gate.nx` source
