`nx_nn_train.nx` source

1// nx_nn_train.nx -- sovereign FIXED-POINT TRAINING (batch gradient descent), the final piece of the neural-audio
2// moonshot machinery. Together with nx_nn (the forward pass) this makes a net LEARNABLE entirely in integer arithmetic:
3// no FPU, no GPU, no third-party autograd -- so a learned PLC / RVQ-codebook predictor can be trained sovereignly on
4// affordable hardware. MSE loss, analytic gradient of a linear layer (dL/dW = (pred - target) . input), Q-scale step.
5// lr_shift sets the learning rate (bigger = smaller, more stable steps). Composes the same matmul as nx_nn. license_tier: ORIGINAL
6
7// one BATCH GD step on a linear layer W[O x I] over T samples X[T x I], targets Y[T x O]. Q-scale fixed-point.
8// Qshift = the pred descale (W's Q-scale plus X's, e.g. W in Q16 . X in Q8 -> pred = acc>>16, so Qshift here would be
9// X's Q... actually pred = (W.X)>>Qshift where Qshift = W's fractional bits). update_shift = the learning rate (W -=
10// grad>>update_shift, applied in W's OWN Q-scale). Keeping W in high precision (Q16) lets the sub-unit GD steps
11// ACCUMULATE -- the fix for the fixed-point "step rounds to 0" trap. grad[O*I] + pred[O] are caller scratch.
12func nnt_step(X: *i64, Y: *i64, T: i64, W: *i64, I: i64, O: i64, Qshift: i64, update_shift: i64, grad: *i64, pred: *i64) -> i64 {
  var g: i64 = 0
  while g < O*I { grad[g] = 0; g = g + 1 }
  var t: i64 = 0
  while t < T {
      var o: i64 = 0
      while o < O {
          var acc: i64 = 0
          var i: i64 = 0
          while i < I { acc = acc + W[o*I + i] * X[t*I + i]; i = i + 1 }
          pred[o] = acc >> Qshift
          o = o + 1
      }
      o = 0
      while o < O {
          let err: i64 = pred[o] - Y[t*O + o]
          var i: i64 = 0
          while i < I { grad[o*I + i] = grad[o*I + i] + err * X[t*I + i]; i = i + 1 }
          o = o + 1
      }
      t = t + 1
  }
  var k: i64 = 0
  while k < O*I { W[k] = W[k] - (grad[k] >> update_shift); k = k + 1 }
  return 0
37}
38// total MSE loss (sum of squared errors) of the linear layer over the T samples
39func nnt_loss(X: *i64, Y: *i64, T: i64, W: *i64, I: i64, O: i64, Qshift: i64) -> i64 {
  var s: i64 = 0
  var t: i64 = 0
  while t < T {
      var o: i64 = 0
      while o < O {
          var acc: i64 = 0
          var i: i64 = 0
          while i < I { acc = acc + W[o*I + i] * X[t*I + i]; i = i + 1 }
          let e: i64 = (acc >> Qshift) - Y[t*O + o]
          s = s + e*e
          o = o + 1
      }
      t = t + 1
  }
  return s
55}
nx_nn_train.nx source

`nx_nn_train.nx` source
