Native training convergence tests (#174) - #180
Merged
ausimian merged 10 commits intoJun 6, 2026
Merged
Conversation
Add a `mode_test/2` macro to `Emily.ConformanceHelper` that emits each conformance test in three lanes from a single body, varying only the bound `predict_opts`: * eval — `[]`, the evaluator path (unchanged behaviour) * native — `native: true, native_fallback: :raise` * fusion — additionally `native_compiled: true` (mx::compile) The native and fusion lanes are tagged `:native` and `:native_compiled` on top of the module's `:conformance`, so `--only conformance` runs all three while `--only native` / `--only native_compiled` select one lane. Both compiler lanes resolve the same HuggingFace repos as the evaluator lane, so the checkpoint download is amortised across them. Convert the six tiny-random suites (DistilBERT, ViT, Whisper, ModernBert, NomicBert, SmolLm3). `Axon.predict`-based tests thread `predict_opts` into the call; `Axon.build`-based smoke tests build `init_fn` on the evaluator and only `predict_fn` under `predict_opts`, so the native gate covers the forward pass rather than random param init. The serving and fast-kernels tests stay eval-only. `native_fallback: :raise` makes the native lanes a no-fallback op-coverage gate: every architecture lowers fully, and the fusion lane holds at the existing 1e-4 reference tolerance (mx::compile only reassociates f32, so its drift is far below the approximate mx::fast::* kernels).
Extend `mode_test/2` to `mode_test/3` with a `lane_tags:` option and run the `*_full` forward-pass conformance suites (ViT-Base, Whisper-tiny) through the native and fusion compilers alongside the evaluator. The full suites pass `lane_tags: false`, so their native and fusion lanes are emitted without the cross-cutting `:native` / `:native_compiled` tags and stay gated behind the suite's own `:vit_full` / `:whisper_full` moduletag. `--only vit_full` now runs all three lanes on the full checkpoint, while `--only native` stays tiny-random only and never pulls a full-size download. Both full forwards lower fully under `native_fallback: :raise` and hold at the pinned 1e-4 reference tolerance in the fusion lane too — `mx::compile` only reassociates f32, so its drift stays well below the approximate `mx::fast::*` kernels (which need 1e-3 ViT / 1e-2 Whisper in the `:fast_kernels_full` variants). The serving and fast-kernels tests remain eval-only.
`Nx.take_along_axis` (the `Nx.Block.TakeAlongAxis` block) now lowers under `compiler: Emily.Compiler, native: true` instead of raising. The MLX op and eager NIF already existed (`Emily.Native.take_along_axis` → `mlx::core::take_along_axis`, used by `Emily.Backend.native_take_along_axis/4`); only the IR path was missing. - ir.ex: add opcode `take_along_axis` (81) and a `lower_block/5` clause mirroring the eager twin — cast indices to s32, emit, coerce to output type — alongside the existing `Nx.Block.Take` lowering. - opcodes.hpp: add `TakeAlongAxis = 81`, bump `kOpcodeCount`, and a dispatch case calling `mx::take_along_axis(in[0], in[1], axis, s)`. - compiler_equivalence_test.exs: native-vs-evaluator bit-identical cases (last axis, axis 0, a 3-D transformer-shaped gather). This was the last op forcing a fallback in `Bumblebee.Text.question_answering`'s answer-span gather, so the DistilBERT-QA `Nx.Serving` forward now compiles fully native and fused. Wire it into the conformance suite: extend `mode_test` with a `:tag` option (stamp every lane) so the serving test runs its eval / native / native_compiled lanes gated behind `:distilbert_full`.
The opcode wire values are hand-maintained in two places — `Emily.IR`'s `@opcodes` map and the `Opcode` enum + `kOpcodeCount` in `c_src/emily/opcodes.hpp` — with nothing enforcing they agree. A mismatch compiles fine and only misbehaves at runtime. Add `Emily.OpcodeParityTest`, which parses the header and asserts both sides are a unique, gap-free `0..N-1` with N == `kOpcodeCount`. The check is value-based, not name-based, so it doesn't depend on the snake_case names matching the PascalCase enum (`negate`/`Negative`, `fast_rms_norm`/`FastRMSNorm`); a name/value permutation that keeps both contiguous is still caught by the equivalence suite. Expose `Emily.IR.opcodes/0` for the test to read the map.
The example documents are Livebooks (`.livemd`), so name the directory to match. Rename `notebooks/` → `livebooks/` (all 8 files) and update every current-state reference: - mix.exs: the `extras:` paths and the `groups_for_extras` regex, plus the nav-group label `Notebooks:` → `Livebooks:`. - scripts/test-livebooks.sh: `NB_DIR` and the path comment. - README.md / ROADMAP.md: the `livebooks/` path and the HexDocs "Livebooks" section references. CHANGELOG.md and PLAN.md keep their `notebooks/` mentions — those are point-in-time historical records, not current-structure docs.
The window-reduce family (`window_sum`/`window_max`/`window_min`/ `window_product` — average and max pooling) now lowers under `compiler: Emily.Compiler, native: true` instead of raising. Unlike a single MLX op, these are composed (pad -> as_strided sliding view -> reduce), so the cores had to be made shareable first. - op_cores.hpp: hoist `contiguous_strides`/`do_pad`/`sliding_windows_view` out of pooling.cpp's anonymous namespace into `emily::ops`, plus a `WindowReduceKind` enum and `window_reduce_core`. pooling.cpp's eager NIFs now call these, so eager and compiled share one implementation. - opcodes.hpp: add WindowSum/Max/Min/Product (82-85), bump kOpcodeCount, dispatch via `window_reduce_core`. - ir.ex: lower_op clause for the four ops, mirroring `Emily.Backend.apply_window_reduce/6` — operands [input, init], iattrs [window, strides, pad_lo, pad_hi, dilations]. The dtype identity (0/1/±inf or int min/max) is baked as a const operand. - equivalence tests: native-vs-eval bit-identical for 2x2 max/sum/min pooling, padded maxpool, and window_product. A small-CNN forward (conv + maxpool) now lowers fully native under `native_fallback: :raise`, unblocking CNN/avg-pool inference. The maxpool backward (`window_scatter_max`) is still pending for CNN training.
Finish the pooling op family so a small-CNN training step lowers with no fallback: the MaxPool/MinPool backward (`window_scatter_max`/ `window_scatter_min`) and `reverse` (the conv-backward kernel flip). - op_cores.hpp: hoist `window_scatter_core` (the select-and-scatter body) out of pooling.cpp's anonymous namespace into `emily::ops`; the eager NIFs now call it, sharing one implementation with the compiled replay. - opcodes.hpp: add WindowScatterMax/Min (86-87), bump kOpcodeCount to 88, dispatch via `window_scatter_core`. - ir.ex: lower_op for `:window_scatter_max`/`:window_scatter_min` (operands [input, source, init]; iattrs [window, strides, pad_lo, pad_hi]), mirroring `Emily.Backend.apply_window_scatter/7`; and `:reverse`, lowered to one `flip` per axis (reusing the existing flip_core / Flip opcode — no new opcode). - equivalence tests: window_scatter_max/min, grad(window_max) (the maxpool backward path), and reverse — all native-vs-eval bit-identical. A conv + maxpool CNN training step (forward + backward + grad + SGD) now lowers fully native under `native_fallback: :raise`, with a loss bit-identical to the evaluator.
Window reductions with `window_dilations > 1` silently returned wrong
values for windows past the first stride positions, on both the eager
backend and the native Expr compiler (they share `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 (overlapping strides). MLX's reduction then takes its
`GeneralStridedReduce` fast path, which assumes a dense, non-overlapping
layout and reads `product(shape)` contiguous elements — over-running the
aliased buffer. The symptom was data-dependent: large dilated cases
happened to over-read valid in-buffer data, while a tiny `{1, 8}` input
read past the allocation (deterministic 0.0 eager, garbage native).
Materialise the view with `mx::contiguous` before the reduce when any
dilation > 1; the general copy reads element-by-element via the real
strides (always in-bounds), yielding a dense buffer the reducer can
safely fast-path. The common non-dilated pooling path is unchanged and
stays copy-free.
Add regression tests: small-tensor dilated sum/max/min/product vs the
BinaryBackend oracle, and a native-vs-evaluator dilated equivalence case
(restoring the coverage dropped in e32b7cc).
Close the training-coverage gap from issue #174: CNN training lowered fully native (window forward/backward, reverse, grad, SGD) but was only verified-lowering, never convergence-tested through the single-NIF replay. Every existing training/* test runs eval-only. Add cnn_native_curve_test and mlp_native_curve_test: each drives the full handwritten training step (forward + backward + grad + SGD) through `compiler: Emily.Compiler, native: true, native_fallback: :raise` for 30/50 steps. The `:raise` makes the run a no-fallback gate — any op that fails to lower raises instead of silently degrading. Each native loss trajectory is checked against the op-by-op evaluator (bit-identical, same MLX kernels) and a BinaryBackend oracle (f32 tolerance), plus a convergence assertion that the loss decreased. Generalize TrainingHelper.run_steps/5 to accept a full jit opts list (or a bare compiler module, as before) so the native lane can pass `native:`/`native_fallback:` through, which a bare `compiler:` can't carry. All existing eval-lane callers are unaffected. Closes #174
Close the Axon-driven half of issue #174: train the LeNet CNN and the dense MLP on real MNIST entirely through the native single-NIF compiler, asserting the same >97% / >96% accuracy as the eval canaries. Axon.Loop.run forwards per-call jit options (it pops only :jit_compile? and :force_garbage_collection?), so passing native: true, native_fallback: :raise reaches the defn jit — confirmed by source and the "Forwarding options" debug log. The issue's "may not forward" caveat does not hold. Under :raise the run is self-proving: once native: true reaches Emily.Compiler it either lowers the whole step to one program or reraises on an un-lowerable op; there is no silent eval-fallback path (that exists only under native_fallback: :eval). So reaching the accuracy assertion proves the forward, cross-entropy loss, backward (window_scatter_max, reverse), and Adam update all lowered fully native. Evaluation runs native too, so the gating accuracy is itself produced by the single-NIF forward. MnistHelper.evaluate gains an optional predict_opts arg (defaults to the eval lane; existing callers unaffected). Both tests are :training_full (download + multi-minute); verified locally — 1 epoch already clears 97%, 5 epochs pass in ~39s. Refs #174
ausimian
changed the base branch from
feat/expr-compiler-window-ops
to
feat/expr-compiler
June 6, 2026 05:37
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Closes #174 — CNN/MLP training lowered fully native, but was only
verified-lowering, never convergence-tested through the single-NIF replay.
test: add native single-NIF training convergence tests—cnn_native_curve_test/mlp_native_curve_testdrive the fullhandwritten training step (forward + backward + grad + SGD) through
native: true, native_fallback: :raisefor 30/50 steps. Each nativeloss trajectory is checked bit-for-bit against the op-by-op evaluator
and to f32 tolerance against a
BinaryBackendoracle, plus aconvergence assertion. Generalises
TrainingHelper.run_steps/5to takea full jit opts list (back-compatible).
test: add native MNIST training convergence canaries (Axon)— aLeNet CNN and a dense MLP train on real MNIST end-to-end through the
native compiler (
:training_full), reaching the same >97% / >96%accuracy as the eval canaries.
Axon.Loop.runforwardsnative:/native_fallback:to the defn jit, so under:raiseacompleted run proves the whole step (forward, cross-entropy, backward,
Adam) lowered with zero fallback.
MnistHelper.evaluategains anoptional
predict_optsarg so evaluation runs native too.Test-only; depends on the window-op lowering in the parent PR.
PR 5/5 (top of stack). Stacked on
…-window-ops; lands infeat/expr-compiler.