Add C weight initializers, LR schedules, and extra losses

CMakeLists.txt

@@ -19,6 +19,9 @@ add_library(nir_c_runtime
   src/fixed_event.c
   src/train.c
   src/stdp_graph.c
+  src/init.c
+  src/schedule.c
+  src/loss.c
   src/rng.c
   src/encode.c
   src/data.c

include/nir_c_runtime/init.h

@@ -0,0 +1,77 @@
+/* Weight initializers for from-scratch C training (#45).
+ *
+ * Every weight in a C-loaded graph otherwise comes from a
+ * Python-exported graph JSON; these fill a fresh buffer so a
+ * zero-Python training run can start somewhere sane, then land
+ * the buffer on a node via `nir_graph_set_weights`.
+ *
+ * Formulas mirror torch.nn.init exactly:
+ *
+ *   uniform:        U(low, high)
+ *   xavier_uniform: U(-b, b), b = gain * sqrt(6 / (fan_in + fan_out))
+ *   xavier_normal:  N(0, s),  s = gain * sqrt(2 / (fan_in + fan_out))
+ *   he_uniform:     U(-b, b), b = gain * sqrt(3 / fan_in)
+ *   he_normal:      N(0, s),  s = gain / sqrt(fan_in)
+ *
+ * Fan conventions (torch `_calculate_fan_in_and_fan_out`, which
+ * `nir_init_fans` implements): a Linear/Affine weight laid out
+ * [out_features, in_features] has fan_in = in_features and
+ * fan_out = out_features; a convolution weight laid out
+ * [out_channels, in_channels, *kernel] multiplies both by the
+ * kernel's receptive-field size (kh * kw for Conv2d, k for
+ * Conv1d). He follows torch's kaiming defaults -- mode 'fan_in'
+ * (only fan_in enters) with an explicit gain, so ReLU nets pass
+ * sqrt(2) (torch `calculate_gain('relu')`, equal to leaky_relu
+ * with a = 0) and Xavier nets pass 1. Use `nir_init_fans` on the
+ * node's weight shape rather than hand-counting.
+ *
+ * All draws come from `nir_rng_t`, so a seeded stream replays
+ * bit-identically. Uniform draws consume one u64 each; normal
+ * draws use Box-Muller pairs (two u64 per pair, the spare
+ * discarded when `count` is odd), so the stream position after
+ * a call depends on `count` -- reseed to replay a prefix.
+ */
+#ifndef NIR_C_RUNTIME_INIT_H
+#define NIR_C_RUNTIME_INIT_H
+
+#include <stddef.h>
+
+#include "nir_c_runtime/rng.h"
+
+/* Torch fan-in/fan-out for a weight shape: shape[0] is output
+ * maps, shape[1] input maps, shape[2..] the kernel. Needs
+ * ndim >= 2, no zero extents, and products that fit size_t. */
+int nir_init_fans(
+    const size_t *shape, size_t ndim, size_t *fan_in,
+    size_t *fan_out
+);
+
+/* Fill `out` with U(low, high). Needs low < high, both finite. */
+int nir_init_uniform(
+    nir_rng_t *rng, float *out, size_t count, float low,
+    float high
+);
+
+/* Xavier/Glorot initializers. Needs both fans nonzero and a
+ * finite gain >= 0 (gain 0 yields all zeros). */
+int nir_init_xavier_uniform(
+    nir_rng_t *rng, float *out, size_t count, size_t fan_in,
+    size_t fan_out, float gain
+);
+int nir_init_xavier_normal(
+    nir_rng_t *rng, float *out, size_t count, size_t fan_in,
+    size_t fan_out, float gain
+);
+
+/* He initializers (torch kaiming_* with mode 'fan_in'). Needs
+ * fan_in nonzero and a finite gain >= 0. */
+int nir_init_he_uniform(
+    nir_rng_t *rng, float *out, size_t count, size_t fan_in,
+    float gain
+);
+int nir_init_he_normal(
+    nir_rng_t *rng, float *out, size_t count, size_t fan_in,
+    float gain
+);
+
+#endif /* NIR_C_RUNTIME_INIT_H */

include/nir_c_runtime/loss.h

@@ -0,0 +1,56 @@
+/* Loss functions beyond single-sample cross-entropy (#47).
+ *
+ * `nir_cross_entropy` (train.h) stays the classifier loss for
+ * logits; these cover the other two common SNN recipes:
+ *
+ *   nir_mse_loss: mean-squared error with 'mean' reduction
+ *   (loss = mean((pred - target)^2), grad = 2(pred - target)/n),
+ *   mirroring torch F.mse_loss on one sample. Serves
+ *   readout/membrane regression -- predicting continuous
+ *   targets such as a decoded intensity or a target membrane
+ *   trace -- where classification-style softmax is wrong.
+ *
+ *   nir_spike_count_loss: softmax cross-entropy over the
+ *   per-class spike-count vector, mirroring torch
+ *   F.cross_entropy(counts, target) on one sample. This is the
+ *   decided spike-count form (cross-entropy over counts, not
+ *   MSE on counts): it matches the snnTorch recipe of summing
+ *   output spikes over the window and classifying the totals,
+ *   and spikeforge's own averaged-readout-plus-cross_entropy
+ *   training path is the same family. The caller accumulates
+ *   the T-step window into `counts` (a plain sum; each step's
+ *   spikes are 0/1, so counts are non-negative by
+ *   construction) and backpropagates through the accumulation
+ *   itself -- every timestep's spikes share their class's
+ *   count gradient -- until the BPTT tape (#39) owns that.
+ *
+ * Both keep the single-sample-first convention: 1-D vectors,
+ * mean reduction for MSE, batch reduction left to the
+ * mini-batch issue (#42).
+ */
+#ifndef NIR_C_RUNTIME_LOSS_H
+#define NIR_C_RUNTIME_LOSS_H
+
+#include <stddef.h>
+
+#include "nir_c_runtime/tensor.h"
+
+/* MSE on one sample. `pred`/`target` must be 1-D with equal
+ * nonzero size and finite data; `grad` needs capacity for the
+ * vector and follows the cross-entropy aliasing rule (exact
+ * aliasing with either input is fine, partial overlap fails). */
+int nir_mse_loss(
+    const nir_tensor_t *pred, const nir_tensor_t *target,
+    float *loss, nir_tensor_t *grad
+);
+
+/* Softmax cross-entropy over spike counts: loss =
+ * -log(softmax(counts)[target]), grad = softmax(counts) with
+ * grad[target] -= 1. Same shape/aliasing rules as
+ * `nir_cross_entropy`. */
+int nir_spike_count_loss(
+    const nir_tensor_t *counts, size_t target, float *loss,
+    nir_tensor_t *grad
+);
+
+#endif /* NIR_C_RUNTIME_LOSS_H */

include/nir_c_runtime/schedule.h

@@ -0,0 +1,52 @@
+/* Learning-rate schedules for the SGD/Adam optimizers (#46).
+ *
+ * Stateless functions computing lr from a step counter, applied
+ * by writing the optimizer's `lr` field before stepping:
+ *
+ *   float lr = 0.0f;
+ *   nir_lr_step_decay(0.1f, 0.5f, 10, meta.step, &lr);
+ *   sgd.lr = lr;
+ *   nir_sgd_step(&sgd, weights, grad, count);
+ *
+ * The same write works for `nir_adam_t.lr`. There is no
+ * scheduler struct to persist: the counter IS the checkpoint's
+ * `step` (completed optimizer steps, 0 for a fresh run), so a
+ * resumed run passes the loaded `meta.step` and continues the
+ * curve exactly where it stopped.
+ *
+ * Formulas mirror torch.optim.lr_scheduler with last_epoch set
+ * to `step` (step 0 always yields base_lr):
+ *
+ *   step_decay:  base_lr * gamma^(step / step_size)
+ *   cosine:      eta_min + (base_lr - eta_min) *
+ *                (1 + cos(pi * step / t_max)) / 2
+ *   exponential: base_lr * gamma^step
+ *
+ * Cosine follows torch past t_max too: beyond the minimum the
+ * curve rises again rather than clamping, so callers that want
+ * a floor should stop the run (or the schedule) at t_max.
+ */
+#ifndef NIR_C_RUNTIME_SCHEDULE_H
+#define NIR_C_RUNTIME_SCHEDULE_H
+
+/* Step decay (torch StepLR). Needs base_lr and gamma finite
+ * and >= 0, step_size >= 1. */
+int nir_lr_step_decay(
+    float base_lr, float gamma, unsigned long step_size,
+    unsigned long step, float *lr
+);
+
+/* Cosine annealing (torch CosineAnnealingLR). Needs base_lr
+ * and eta_min finite and >= 0, t_max >= 1. */
+int nir_lr_cosine(
+    float base_lr, float eta_min, unsigned long t_max,
+    unsigned long step, float *lr
+);
+
+/* Exponential decay (torch ExponentialLR). Needs base_lr and
+ * gamma finite and >= 0. */
+int nir_lr_exponential(
+    float base_lr, float gamma, unsigned long step, float *lr
+);
+
+#endif /* NIR_C_RUNTIME_SCHEDULE_H */

src/init.c

@@ -0,0 +1,137 @@
+#include "nir_c_runtime/init.h"
+
+#include <math.h>
+#include <stdint.h>
+
+static int field_size(
+    const size_t *shape, size_t ndim, size_t *field
+) {
+    size_t f = 1;
+    for (size_t i = 2; i < ndim; i++) {
+        if (shape[i] == 0 || f > SIZE_MAX / shape[i]) {
+            return 0;
+        }
+        f *= shape[i];
+    }
+    *field = f;
+    return 1;
+}
+
+int nir_init_fans(
+    const size_t *shape, size_t ndim, size_t *fan_in,
+    size_t *fan_out
+) {
+    size_t field = 0;
+    if (shape == NULL || ndim < 2 || fan_in == NULL ||
+        fan_out == NULL || shape[0] == 0 || shape[1] == 0 ||
+        !field_size(shape, ndim, &field) ||
+        shape[1] > SIZE_MAX / field ||
+        shape[0] > SIZE_MAX / field) {
+        return 0;
+    }
+    *fan_in = shape[1] * field;
+    *fan_out = shape[0] * field;
+    return 1;
+}
+
+static int check_fill(
+    const nir_rng_t *rng, const float *out, size_t count
+) {
+    return rng != NULL && out != NULL && count > 0;
+}
+
+int nir_init_uniform(
+    nir_rng_t *rng, float *out, size_t count, float low,
+    float high
+) {
+    if (!check_fill(rng, out, count) || !isfinite(low) ||
+        !isfinite(high) || !(low < high)) {
+        return 0;
+    }
+    double span = (double)high - (double)low;
+    for (size_t i = 0; i < count; i++) {
+        double u = nir_rng_next_f64(rng);
+        out[i] = (float)((double)low + span * u);
+    }
+    return 1;
+}
+
+/* One standard-normal draw via Box-Muller. u1 == 0 would feed
+ * log(0), so resample it; the stream stays deterministic. */
+static double gaussian(nir_rng_t *rng) {
+    double u1 = 0.0;
+    do {
+        u1 = nir_rng_next_f64(rng);
+    } while (u1 == 0.0);
+    double u2 = nir_rng_next_f64(rng);
+    double r = sqrt(-2.0 * log(u1));
+    return r * cos(2.0 * acos(-1.0) * u2);
+}
+
+static int check_gain(float gain) {
+    return isfinite(gain) && gain >= 0.0f;
+}
+
+int nir_init_xavier_uniform(
+    nir_rng_t *rng, float *out, size_t count, size_t fan_in,
+    size_t fan_out, float gain
+) {
+    if (!check_fill(rng, out, count) || fan_in == 0 ||
+        fan_out == 0 || !check_gain(gain)) {
+        return 0;
+    }
+    double total = (double)fan_in + (double)fan_out;
+    double bound = (double)gain * sqrt(6.0 / total);
+    for (size_t i = 0; i < count; i++) {
+        double u = nir_rng_next_f64(rng);
+        out[i] = (float)(-bound + 2.0 * bound * u);
+    }
+    return 1;
+}
+
+int nir_init_xavier_normal(
+    nir_rng_t *rng, float *out, size_t count, size_t fan_in,
+    size_t fan_out, float gain
+) {
+    if (!check_fill(rng, out, count) || fan_in == 0 ||
+        fan_out == 0 || !check_gain(gain)) {
+        return 0;
+    }
+    double total = (double)fan_in + (double)fan_out;
+    double std = (double)gain * sqrt(2.0 / total);
+    for (size_t i = 0; i < count; i++) {
+        out[i] = (float)(std * gaussian(rng));
+    }
+    return 1;
+}
+
+int nir_init_he_uniform(
+    nir_rng_t *rng, float *out, size_t count, size_t fan_in,
+    float gain
+) {
+    if (!check_fill(rng, out, count) || fan_in == 0 ||
+        !check_gain(gain)) {
+        return 0;
+    }
+    double bound = (double)gain * sqrt(3.0 / (double)fan_in);
+    for (size_t i = 0; i < count; i++) {
+        double u = nir_rng_next_f64(rng);
+        out[i] = (float)(-bound + 2.0 * bound * u);
+    }
+    return 1;
+}
+
+int nir_init_he_normal(
+    nir_rng_t *rng, float *out, size_t count, size_t fan_in,
+    float gain
+) {
+    if (!check_fill(rng, out, count) || fan_in == 0 ||
+        !check_gain(gain)) {
+        return 0;
+    }
+    double std = (double)gain / sqrt((double)fan_in);
+    for (size_t i = 0; i < count; i++) {
+        out[i] = (float)(std * gaussian(rng));
+    }
+    return 1;
+}