Add mini-batch accumulation, mean reduction, and seeded shuffling (closes #42)

CMakeLists.txt

@@ -26,6 +26,7 @@ add_library(nir_c_runtime
   src/encode.c
   src/data.c
   src/checkpoint.c
+  src/minibatch.c
 )
 target_include_directories(nir_c_runtime PUBLIC include)
 target_link_libraries(nir_c_runtime PUBLIC m)

include/nir_c_runtime/data.h

@@ -8,7 +8,9 @@
  * 1/255 and the Normalize is an identity -- so C decodes each
  * pixel as v / 255.0f. Multi-byte samples land in `nir_tensor_t`
  * frames (batch-1, one tensor per sample); batching stays the
- * caller's loop until the tensor ABI grows a batch axis.
+ * caller's loop with the mean-reduction convention in
+ * minibatch.h, which also records the #42 decision against a
+ * batch axis on the tensor ABI.
  */
 #ifndef NIR_C_RUNTIME_DATA_H
 #define NIR_C_RUNTIME_DATA_H

include/nir_c_runtime/loss.h

@@ -25,8 +25,8 @@
  *   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).
+ * mean reduction for MSE, batch reduction via the mean
+ * convention in minibatch.h.
  */
 #ifndef NIR_C_RUNTIME_LOSS_H
 #define NIR_C_RUNTIME_LOSS_H

include/nir_c_runtime/minibatch.h

@@ -0,0 +1,80 @@
+/* Mini-batch training over batch-1 tensors (#42).
+ *
+ * Reduction convention (this is the documented contract; train.h
+ * and data.h point here): the caller loops over one sample at a
+ * time and reduces with a mean, matching torch's default
+ * `reduction='mean'`:
+ *
+ *   batch loss = mean(sample losses)
+ *   batch grad = mean(sample grads)
+ *
+ * Per batch the caller zeroes one accumulator per parameter
+ * buffer (`nir_minibatch_zero`), adds each sample's gradient
+ * (`nir_minibatch_accumulate`), scales the sums by 1/B
+ * (`nir_minibatch_scale`), and only then takes the optimizer
+ * step. Losses reduce the same way through
+ * `nir_minibatch_mean_loss`. The optimizers in train.h always
+ * see the mean gradient -- scale exactly once per batch.
+ *
+ * Within one sample the single-sample losses keep their own
+ * rules: cross-entropy (train.h) and spike-count loss (loss.h)
+ * are per-logit-vector, MSE (loss.h) already means over its
+ * elements. Those per-sample values are what the batch mean
+ * averages; nothing here re-weights them.
+ *
+ * Batch-axis decision (#42, decided): no batch axis on
+ * `nir_tensor_t`. The accumulation loop above keeps the tensor
+ * ABI and every inference call site unchanged, and the extra
+ * adds are exact and cheap next to a forward/backward pass.
+ * Revisit only with measured evidence that the loop -- not the
+ * math under it -- is the bottleneck or the bug farm.
+ *
+ * Epoch shuffling is Fisher-Yates over an index array
+ * (`nir_sampler_arange` + `nir_sampler_shuffle`), driven by the
+ * seeded `nir_rng_t` stream so an epoch order replays exactly
+ * from its seed.
+ */
+#ifndef NIR_C_RUNTIME_MINIBATCH_H
+#define NIR_C_RUNTIME_MINIBATCH_H
+
+#include <stddef.h>
+
+#include "nir_c_runtime/rng.h"
+
+/* Zero an accumulator of `n` floats. `accum` must hold `n`
+ * floats; `n` may be 0. */
+void nir_minibatch_zero(float *accum, size_t n);
+
+/* Add `grad` into `accum` elementwise. Both buffers must hold
+ * `n` floats with `n` nonzero. Fails (return 0, `accum`
+ * untouched) on NULL buffers or non-finite data. */
+int nir_minibatch_accumulate(
+    float *accum, const float *grad, size_t n
+);
+
+/* Scale `accum` by `scale` elementwise -- the 1/B step that
+ * turns accumulated sums into the mean gradient. Fails
+ * (return 0, `accum` untouched) on a NULL buffer, `n` of 0,
+ * or non-finite scale or data. */
+int nir_minibatch_scale(float *accum, size_t n, float scale);
+
+/* Mean of `n` per-sample losses, accumulated in float64 and
+ * rounded once. Fails (return 0) on NULL buffers, `n` of 0,
+ * or a non-finite entry. */
+int nir_minibatch_mean_loss(
+    const float *losses, size_t n, float *mean
+);
+
+/* Fill `indices` with 0..n-1. `indices` must hold `n`
+ * entries; `n` may be 0. */
+void nir_sampler_arange(size_t *indices, size_t n);
+
+/* Fisher-Yates shuffle of `indices` in place, drawing one
+ * `nir_rng_next_u64` per position from `i = n-1` down to 1
+ * with `j = draw % (i + 1)`. A seed therefore replays the
+ * same permutation on any platform. Fails (return 0) on a
+ * NULL `rng`, or a NULL `indices` with `n` nonzero; `n` of
+ * 0 or 1 is a successful no-op. */
+int nir_sampler_shuffle(nir_rng_t *rng, size_t *indices, size_t n);
+
+#endif /* NIR_C_RUNTIME_MINIBATCH_H */

include/nir_c_runtime/train.h

@@ -28,7 +28,8 @@ int nir_surrogate_from_name(
 /* Softmax cross-entropy for one logit vector, matching
  * torch.nn.functional.cross_entropy on a single sample: loss =
  * -log(softmax(logits)[target]), grad = softmax(logits) with
- * grad[target] -= 1. Batch reduction is the caller's job. */
+ * grad[target] -= 1. Batch reduction is the caller's loop with
+ * the mean convention in minibatch.h. */
 int nir_cross_entropy(
     const nir_tensor_t *logits, size_t target, float *loss,
     nir_tensor_t *grad

src/minibatch.c

@@ -0,0 +1,75 @@
+#include "nir_c_runtime/minibatch.h"
+
+#include <math.h>
+
+void nir_minibatch_zero(float *accum, size_t n) {
+    for (size_t i = 0; i < n; i++) {
+        accum[i] = 0.0f;
+    }
+}
+
+static int check_finite(const float *data, size_t n) {
+    for (size_t i = 0; i < n; i++) {
+        if (!isfinite(data[i])) return 0;
+    }
+    return 1;
+}
+
+int nir_minibatch_accumulate(
+    float *accum, const float *grad, size_t n
+) {
+    if (accum == NULL || grad == NULL || n == 0 ||
+        !check_finite(accum, n) || !check_finite(grad, n)) {
+        return 0;
+    }
+    for (size_t i = 0; i < n; i++) {
+        accum[i] += grad[i];
+    }
+    return 1;
+}
+
+int nir_minibatch_scale(float *accum, size_t n, float scale) {
+    if (accum == NULL || n == 0 || !isfinite(scale) ||
+        !check_finite(accum, n)) {
+        return 0;
+    }
+    for (size_t i = 0; i < n; i++) {
+        accum[i] *= scale;
+    }
+    return 1;
+}
+
+int nir_minibatch_mean_loss(
+    const float *losses, size_t n, float *mean
+) {
+    if (losses == NULL || mean == NULL || n == 0 ||
+        !check_finite(losses, n)) {
+        return 0;
+    }
+    double total = 0.0;
+    for (size_t i = 0; i < n; i++) {
+        total += (double)losses[i];
+    }
+    *mean = (float)(total / (double)n);
+    return 1;
+}
+
+void nir_sampler_arange(size_t *indices, size_t n) {
+    for (size_t i = 0; i < n; i++) {
+        indices[i] = i;
+    }
+}
+
+int nir_sampler_shuffle(nir_rng_t *rng, size_t *indices, size_t n) {
+    if (rng == NULL || (indices == NULL && n > 0)) {
+        return 0;
+    }
+    for (size_t i = n; i > 1; i--) {
+        size_t j = (size_t)(nir_rng_next_u64(rng) % i);
+        size_t last = i - 1;
+        size_t tmp = indices[last];
+        indices[last] = indices[j];
+        indices[j] = tmp;
+    }
+    return 1;
+}

tests/CMakeLists.txt

@@ -26,6 +26,7 @@ nir_add_test(test_rng)
 nir_add_test(test_encode)
 nir_add_test(test_data)
 nir_add_test(test_checkpoint)
+nir_add_test(test_minibatch)
 
 # OOM sweep includes src/json.c directly, so it must not link the library
 # (that would duplicate the json_* symbols); it only needs the headers.