Native training convergence tests (#174)

README.md

@@ -191,16 +191,16 @@ Nx.Defn.global_default_options(compiler: Emily.Compiler)
 Bumblebee inference works with no further configuration once the
 backend is installed — see the conformance suites under
 `test/emily/conformance/` for worked DistilBERT, Qwen3, ViT, and
-Whisper pipelines, and the Notebooks section of the HexDocs nav for
+Whisper pipelines, and the Livebooks section of the HexDocs nav for
 runnable Livebooks.
 
 The low-level tensor API (`Emily.from_binary/3`, `to_binary/1`,
 `shape/1`, `dtype/1`, `eval/1`) remains available for diagnostics
 and direct MLX round-trips, but most users should go through Nx.
 
-## Notebooks
+## Livebooks
 
-End-to-end Livebooks under `notebooks/`. Each one declares its own
+End-to-end Livebooks under `livebooks/`. Each one declares its own
 `Mix.install/2` block and pins `Emily.Backend` as the default Nx
 backend, so they're self-contained — open in Livebook and run.
 
@@ -339,8 +339,8 @@ caller-facing API; the only difference is whether the calling
 process wraps the call in `Emily.Stream.with_stream/2` and whether
 you run one serving or many.
 
-See `Emily.Stream` for the API and the `qwen3_quantized` notebook
-under Notebooks for a worked multi-stream example.
+See `Emily.Stream` for the API and the `qwen3_quantized` livebook
+under Livebooks for a worked multi-stream example.
 
 ## Observability
 
@@ -438,7 +438,7 @@ be introduced in the layer where its test fails.
 ## Documentation
 
   * [HexDocs](https://hexdocs.pm/emily) — per-module API docs and
-    runnable notebooks.
+    runnable livebooks.
   * [`ARCHITECTURE.md`](ARCHITECTURE.md) — current shape of the
     library: layer boundaries, design decisions, concurrency and
     memory model, observability surface.

RELEASE.md

@@ -50,6 +50,34 @@
   selection on the native path. Remaining gaps (`gather`/scatter,
   pooling/`window_*`, cumulative) continue to work via the graceful fallback.
 
+- **`take_along_axis` lowers natively** — `Nx.take_along_axis` (the
+  `Nx.Block.TakeAlongAxis` block) now compiles under the native single-NIF
+  path, mirroring `Emily.Backend.native_take_along_axis/4` (cast indices to
+  s32, then `mlx::core::take_along_axis`) bit-for-bit. This was the last op
+  forcing a fallback in `Bumblebee.Text.question_answering`'s answer-span
+  gather, so a DistilBERT question-answering `Nx.Serving` forward now runs
+  fully native — and fused — under `native_fallback: :raise`.
+
+- **Window (pooling) ops lower natively — forward and backward.** The
+  forward window family (`window_sum`/`window_max`/`window_min`/
+  `window_product`, i.e. average and max pooling), the select-and-scatter
+  backward (`window_scatter_max`/`window_scatter_min`, the MaxPool/MinPool
+  gradient), and `reverse` (the conv-backward kernel flip) all now compile
+  under the native single-NIF path instead of falling back. The pad →
+  sliding-window → reduce/scatter cores moved into `emily/op_cores.hpp` so
+  the eager NIFs and the compiled replay share one implementation. A
+  small-CNN **training step** (conv + maxpool forward and backward, grad,
+  SGD) now lowers fully native under `native_fallback: :raise`, producing a
+  loss bit-identical to the evaluator. Native training is now
+  **convergence**-tested, not just verified-lowering: handwritten CNN
+  (30 SGD steps) and MLP (50 steps) trajectories match the op-by-op
+  evaluator bit-for-bit and a `BinaryBackend` oracle to f32 tolerance, and
+  full **Axon** training drives native end-to-end — `Axon.Loop.run`
+  forwards `native: true`/`native_fallback:` to the defn jit, so a LeNet
+  CNN and a dense MLP train on real MNIST entirely through the single-NIF
+  path (forward, categorical-cross-entropy, backward, Adam) and reach the
+  same >97% / >96% accuracy as the evaluator (`:training_full`).
+
 - **`Bumblebee.Text.generation` compiles fully native — greedy and sampling.**
   The headline result: an end-to-end Bumblebee generation (the transformer
   forward, the `defn while` decode loop, dynamic KV-cache writes, `cumsum`
@@ -137,3 +165,15 @@
   tensor-offset `fast_rope/8`. Note: feed the kernel the 4-D
   `{batch, heads, seq, head_dim}` layout — in 3-D, MLX 0.31 mis-rotates
   single-token (`seq == 1`) inputs.
+
+### Fixed
+
+- **Dilated window reductions (`window_dilations > 1`) returned wrong values.**
+  `window_sum`/`window_max`/`window_min`/`window_product` with a dilated kernel
+  silently produced garbage for windows past the first stride positions, on both
+  the eager backend and the native compiler (they share the window-reduce core).
+  A dilated kernel axis gets an `as_strided` stride > 1, so the sliding-window
+  view aliases fewer physical elements than its logical size; MLX's strided-reduce
+  fast path then read past the aliased buffer. The view is now materialised
+  contiguously before the reduce when any dilation > 1 (the common non-dilated
+  pooling path is unchanged and stays copy-free).

ROADMAP.md

@@ -74,5 +74,5 @@ Tracking checklist:
   * `CHANGELOG.md` accumulated across releases (it is, since 0.3.0).
   * `MAINTAINING.md` reflects the precompiled-NIF release flow (it
     does, since 0.3.0).
-  * Worked Bumblebee + quantized-Qwen3 examples in `notebooks/`
-    (present and grouped in the HexDocs Notebooks section).
+  * Worked Bumblebee + quantized-Qwen3 examples in `livebooks/`
+    (present and grouped in the HexDocs Livebooks section).