Merge pull request #18 from ausimian/m9-gradient-training

M9: gradient conformance and training primitives

Author: ausimian

PLAN.md

@@ -21,7 +21,9 @@ checklist so future-us understands the trade-offs.
 
 - Ahead-of-time compilation (IREE-style). Complementary, separate effort.
 - Windows or non-Apple-Silicon Linux GPU. CPU-only Linux is a nice-to-have for CI.
-- Training / gradients beyond what `Nx.Defn` gives for free. Inference is the priority.
+- Training framework features beyond `Nx.Defn.grad`: distributed training,
+  mixed-precision master weights, a native optimizer library. (Autodiff +
+  small-scale training loops: in scope from M9.)
 - Drop-in replacement for EMLX. We borrow where it's clearly right, but
   we're not constrained by its API.
 - `Emily.Stream` as a public API — MLX streams stay internal in v1.
@@ -174,7 +176,7 @@ throughput number.
 
 **Exit:** Axon MLPs forward with `compiler: Emily.Compiler`; results
 match `Nx.Defn.Evaluator` running on the same backend within float
-tolerance. (Training is out of scope for v1.)
+tolerance. (Training via `Nx.Defn.grad` lands in M9.)
 
 ### M6 — `mlx::core::compile` wrapping — **dropped**
 
@@ -254,7 +256,87 @@ exists; only the Backend callback still routes through the
 BinaryBackend fallback). Gated on the M7 ViT and Whisper suites
 staying green through the switchover.
 
-### M9 — 1.0 release
+### M9 — Gradient conformance and training primitives
+
+Training on Emily has been technically possible since M2 —
+`Nx.Defn.grad` is pure symbolic differentiation in Elixir and lowers
+to the same ops the forward pass uses. M9 turns "possible" into
+"usable" by (a) lifting the training-hot indexing ops off the
+`via_binary` fallback and (b) building the test scaffolding needed
+to trust a gradient.
+
+**Primitives.** `Nx.Defn.grad` of indexing-shaped ops lands on
+`indexed_add`; every such backward currently ships to BinaryBackend
+and back. Lift to native MLX:
+
+- `indexed_add` → `mlx::core::scatter_add`
+- `indexed_put` → `mlx::core::scatter`
+- `gather` → `mlx::core::gather`
+
+Window reductions stay on `via_binary` in M9 — pool-based conv
+training is scoped to M10.
+
+**Testing — Layers 4 (Grad) and 5 (Training):**
+
+1. **Grad-equivalence property tests** — for a zoo of `defn`-expressible
+   functions f, assert `Nx.Defn.grad(f)` on `Emily.Backend` matches the
+   same grad on `Nx.BinaryBackend` within dtype-appropriate tolerance.
+   Reuses M2's StreamData harness; the zoo excludes non-differentiable
+   ops (`argmax`, `argmin`, `floor`, `sign`, comparisons).
+2. **Numerical finite-difference oracle** — for the differentiable
+   subset, assert `(f(x+ε) - f(x-ε)) / 2ε ≈ grad(f)(x)`. Tolerance is
+   per-op and documented; f32 central differences bottom out around
+   1e-3 relative, so symbolic-grad tolerance must be relaxed
+   accordingly where this is the oracle. Pilot on 3–4 ops before
+   scaling the harness.
+3. **Training curve-matching** — handwritten MLP and handwritten
+   transformer-block training step, fixed seed, 50–200 steps; assert
+   per-step loss trajectory matches `Nx.BinaryBackend` within
+   tolerance. No Axon dependency in this tier — fewer moving parts
+   when a test goes red.
+4. **Training memory soak** (`test/soak/training_test.exs`,
+   `@tag :soak`) — 1k training steps; MLX memory returns to baseline
+   after `clear_cache/0`. Training exercises a different allocator
+   pattern than inference (param-grad pairs, optimizer state,
+   long-lived activation caches).
+5. **`:training_full`** (opt-in via `--only training_full`, **not**
+   on default CI) — Axon MLP on MNIST → >97% test accuracy. Catches
+   systemic numerical drift that curve-matching misses because both
+   sides use `Nx.BinaryBackend` as the oracle.
+
+Axon is added as a **test-only** dependency, used only by the
+`:training_full` tier.
+
+**Risks specific to this milestone:**
+
+- f32 tolerance calibration for oracle (2) is per-op; the harness
+  must support per-op tolerance tables, not a single global epsilon.
+- Random-key flow through `Emily.Compiler` needs an explicit test:
+  grad through `dropout` with threaded keys, two invocations of the
+  same compiled function must advance the RNG correctly.
+- MLX scatter semantics (out-of-bounds handling, tie-breaking) may
+  differ from Nx expectations. Document divergence; encode property
+  exclusions if needed.
+
+**Exit:** oracles (1)–(3) green in default CI; (4) and (5) green in
+opt-in CI job.
+
+### M10 — Conv-pool training
+
+Lift window reductions (`window_sum`, `window_max`, `window_min`,
+`window_product`, `window_scatter_max`, `window_scatter_min`) off
+`via_binary` onto their native MLX counterparts. This closes the
+last gap in the training primitive set and unblocks pool-based conv
+models (small CNNs, ViT classifier heads trained from scratch).
+
+Scope is narrow: the lifts are mechanical per-op changes. Test
+coverage extends the M9 grad-equivalence and curve-matching zoo to
+cover the new ops, plus a small-CNN MNIST run in `:training_full`.
+
+**Exit:** grad-equivalence on window ops green; small-CNN MNIST
+training converges in `:training_full`.
+
+### M11 — 1.0 release
 
 - API docs, HexDocs, README with a worked Bumblebee example
 - Hex release (public), versioned per conventions (`@version` in mix.exs)
@@ -267,6 +349,8 @@ staying green through the switchover.
 | Native | Hand-computed expected values | ExUnit unit tests |
 | Backend | `Nx.BinaryBackend` on the same inputs | StreamData property tests + Nx conformance |
 | Compiler | `Emily.Backend` in non-defn mode | Equivalence tests (same function, two modes) |
+| Grad | `Nx.BinaryBackend` grad + finite differences | StreamData property tests + numerical oracle |
+| Training | `Nx.BinaryBackend` loss trajectory | Curve-matching; MNIST convergence (`:training_full`, opt-in) |
 | E2E | EXLA-produced golden outputs | Conformance tests with cached weights |
 
 A bug can only be introduced in the layer where its test fails — no
@@ -275,13 +359,19 @@ cross-layer mystery bugs.
 Additional harnesses:
 - **Memory soak** (`test/soak/memory_test.exs`, `@tag :soak`): 10k
   iterations; MLX memory stats asserted to return to baseline.
+- **Training memory soak** (`test/soak/training_test.exs`,
+  `@tag :soak`, from M9): 1k training steps; baseline restored after
+  `clear_cache/0`.
 - **Concurrency** (`test/soak/concurrency_test.exs`, `@tag :soak`):
   parallel inference; determinism + no crashes.
 - **Benchmarks** (`bench/`): Benchee scripts; results logged in
   `RELEASE.md` per version.
 - **Conformance vs EXLA** (CI matrix, Mac for Emily + Linux+CUDA for
   EXLA oracle): same model, same input; runs on every PR touching
   Backend.
+- **Convergence** (`:training_full`, from M9; opt-in CI job): Axon
+  training loop on MNIST; catches numerical drift the curve-matching
+  oracle can't see.
 
 ## Risks and mitigations
 

RELEASE.md

@@ -2,6 +2,91 @@
 
 ## Added
 
+- M9 — Gradient conformance and training primitives. Makes
+  `Nx.Defn.grad` usable on Emily by lifting the three ops that grad
+  lands on most heavily off the `via_binary` fallback, and adds the
+  test scaffolding to trust a gradient. Training on Emily has been
+  technically possible since M2 (grad is symbolic in Elixir and
+  lowers to forward ops), but every embedding-style backward was
+  round-tripping to `Nx.BinaryBackend`.
+  - **Native `Native.gather/4`, `Native.scatter/4`, `Native.scatter_add/4`**
+    backing `mx::gather`, `mx::scatter`, `mx::scatter_add`. List-of-
+    per-axis-indices form (MLX's native shape) rather than the
+    `{N, rank}` tensor Nx passes around — the Backend layer does the
+    translation.
+  - **`Emily.Backend.gather/4`, `indexed_add/5`, `indexed_put/5`**
+    rewired. `gather` retains the single-axis `Native.take` fast path
+    (now with an explicit output reshape — fixes a latent shape lie
+    exposed when downstream MLX ops inspect the ref shape directly,
+    e.g. `Nx.gather` followed by `Nx.dot` in a grad); the multi-axis
+    path is native. Shared `apply_scatter/6` helper handles index
+    splitting (`split_indices_per_axis/3`) and MLX's updates-shape
+    rewrap — Nx ships `{batch ++ non_indexed_dims}` but MLX requires
+    `{batch ++ per_axis_slot}` with a length-1 dim at every indexed-
+    axis position (`updates_shape_for_scatter/3`). Fallback to
+    `via_binary` on shapes outside the covered contract.
+  - **Duplicate-index divergence.** MLX `scatter` is parallel and
+    unordered on duplicate indices; `Nx.indexed_put` is deterministic
+    last-write. Documented in the NIF and Backend; grad property
+    generators dedupe. `scatter_add` is commutative so duplicates
+    accumulate correctly either way.
+  - **`test/emily/native_test.exs`** (+5 cases) — scalar-write,
+    partial-axis slice write, and the load-bearing `{B, L, D}` target
+    with `axes: [0, 1]` rewrap case.
+  - **`test/emily/backend_test.exs`** (+8 tests) — targeted cases for
+    all three ops exercising the Backend translation paths end-to-end
+    via Nx.
+  - **`test/emily/grad/grad_equivalence_test.exs`** — property zoo
+    covering sum, dot, reshape∘transpose, broadcast, gather,
+    indexed_add, plus two composition cases (gather→dot→softmax and
+    a mini-attention block). Each case runs under `compiler:
+    Emily.Compiler` on Emily.Backend and `compiler:
+    Nx.Defn.Evaluator` on `Nx.BinaryBackend`, asserting the
+    grad matches within a grad-scaled tolerance. Also covers the
+    M9 PRNG-key-threading risk called out in `PLAN.md`: grad through
+    a `Nx.Random.uniform_split`-driven dropout with fixed keys
+    produces bit-identical results across repeat runs.
+  - **`test/emily/grad/finite_diff_test.exs`** +
+    **`test/support/grad_helper.ex`** — finite-difference numerical-
+    gradient oracle. Pilot of four ops (`sum`, `dot`, `logsumexp`,
+    `sigmoid`) with per-op tolerance tables; the harness catches the
+    class of bug where symbolic-grad-on-Emily and symbolic-grad-on-
+    BinaryBackend *agree* but are both wrong.
+  - **`test/emily/training/mlp_curve_test.exs`** +
+    **`test/emily/training/transformer_block_curve_test.exs`** +
+    **`test/support/training_helper.ex`** — handwritten MLP and
+    single transformer block (attention + FFN + residuals) trained
+    with vanilla SGD for 50 steps on a fixed synthetic batch. Two
+    tolerance bands asserted in each: per-step loss `rtol = 1e-3`
+    (silent-drift canary — MLX parallel reductions diverge from
+    BinaryBackend sequential reductions, so strict bit-match would
+    be flaky) and final-loss `rtol = 1e-4` (convergence
+    correctness). No Axon dep at this tier; the training loop is
+    handwritten so a red test points at backend/grad numerics.
+  - **`test/soak/training_test.exs`** — 1k-iteration training-loop
+    memory soak, `@moduletag :soak` (default suite). Reuses the
+    handwritten MLP; asserts active memory returns within 2 MB of
+    baseline after `Native.clear_cache/0`. Training exercises a
+    different allocator pattern than inference (param–grad pairs,
+    activations), hence a dedicated soak alongside
+    `memory_test.exs`.
+  - **`test/emily/training/mnist_full_test.exs`** — opt-in MNIST
+    convergence canary, `@moduletag :training_full` (excluded by
+    default; run via `mix test --only training_full`). Loads MNIST
+    through `scidata`, trains an Axon MLP on `Emily.Compiler` for 5
+    epochs, asserts >97% test accuracy. Catches systemic grad drift
+    that curve-matching misses because curve-matching uses
+    BinaryBackend as its own oracle — MNIST convergence is an
+    independent cross-check against real-world training dynamics.
+    Typical wall time ~10 s on Apple Silicon.
+  - **`{:scidata, "~> 0.1", only: :test}`** added for MNIST loading.
+    Kept test-only; Emily itself has no dataset-loading dep.
+  - **`test/test_helper.exs`** — `:training_full` added to the
+    default-exclude list, documented alongside the other heavyweight
+    tags.
+  - **`PLAN.md`** — M9 scope formalized; M10 (conv-pool training)
+    and M11 (1.0 release) renumbered.
+
 - M8 — Native `conv`. `Emily.Backend.conv/4` now dispatches directly
   to `Native.conv_general` (already bound to `mlx::core::conv_general`
   since M1) instead of round-tripping through `Nx.BinaryBackend`.

c_src/ops/index.cpp

@@ -10,7 +10,9 @@
 
 namespace mx = mlx::core;
 using emily::Tensor;
+using emily::to_int_vec;
 using emily::to_mlx_shape;
+using emily::unwrap_all;
 using emily::wrap;
 
 namespace {
@@ -109,4 +111,59 @@ fine::ResourcePtr<Tensor> scatter_add_axis(
 }
 FINE_NIF(scatter_add_axis, 0);
 
+// gather/4 — multi-axis fancy-indexing gather. `indices` is a list of
+// integer index tensors (one per axis in `axes`); all index tensors
+// share a common leading batch shape. `slice_sizes` has length
+// rank(a); entries for the indexed axes are 1 and entries for the
+// remaining axes equal a.shape()[axis] (the window size per gather).
+// Result shape is `batch_shape ++ slice_sizes`.
+fine::ResourcePtr<Tensor> gather(
+    ErlNifEnv *,
+    fine::ResourcePtr<Tensor> a,
+    std::vector<fine::ResourcePtr<Tensor>> indices,
+    std::vector<int64_t> axes,
+    std::vector<int64_t> slice_sizes) {
+  return wrap(mx::gather(
+      a->array,
+      unwrap_all(indices),
+      to_int_vec(axes),
+      to_mlx_shape(slice_sizes)));
+}
+FINE_NIF(gather, 0);
+
+// scatter/4 — multi-axis fancy-indexing scatter (overwrite). MLX
+// requires `updates.ndim() == indices[0].ndim() + a.ndim()`; each
+// update is a slice of a.ndim() dims written at the index site. The
+// Backend layer is responsible for reshaping Nx-shaped updates into
+// MLX's expected shape.
+//
+// Note: MLX scatter with duplicate indices has unordered semantics
+// (parallel write) — it is not last-write-wins like Nx.indexed_put.
+// Callers requiring deterministic duplicate handling must dedupe
+// beforehand.
+fine::ResourcePtr<Tensor> scatter(
+    ErlNifEnv *,
+    fine::ResourcePtr<Tensor> a,
+    std::vector<fine::ResourcePtr<Tensor>> indices,
+    fine::ResourcePtr<Tensor> updates,
+    std::vector<int64_t> axes) {
+  return wrap(mx::scatter(
+      a->array, unwrap_all(indices), updates->array, to_int_vec(axes)));
+}
+FINE_NIF(scatter, 0);
+
+// scatter_add/4 — multi-axis scatter-accumulate. Same shape contract
+// as scatter; duplicate indices accumulate deterministically (add is
+// commutative, so parallel scatter order doesn't affect the result).
+fine::ResourcePtr<Tensor> scatter_add(
+    ErlNifEnv *,
+    fine::ResourcePtr<Tensor> a,
+    std::vector<fine::ResourcePtr<Tensor>> indices,
+    fine::ResourcePtr<Tensor> updates,
+    std::vector<int64_t> axes) {
+  return wrap(mx::scatter_add(
+      a->array, unwrap_all(indices), updates->array, to_int_vec(axes)));
+}
+FINE_NIF(scatter_add, 0);
+
 } // namespace