diff --git a/README.md b/README.md index 46bc0bb..a61a7c8 100644 --- a/README.md +++ b/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. diff --git a/RELEASE.md b/RELEASE.md index 67e309e..6973ede 100644 --- a/RELEASE.md +++ b/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). diff --git a/ROADMAP.md b/ROADMAP.md index 9b16bee..6d77eb6 100644 --- a/ROADMAP.md +++ b/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). diff --git a/c_src/emily/op_cores.hpp b/c_src/emily/op_cores.hpp index 7f0bb5a..fa8a590 100644 --- a/c_src/emily/op_cores.hpp +++ b/c_src/emily/op_cores.hpp @@ -12,8 +12,12 @@ #include +#include +#include +#include #include #include +#include namespace emily::ops { @@ -58,4 +62,302 @@ inline mx::array flip_core(const mx::array &a, int64_t axis, mx::Stream &s) { s); } +// --- Window / pooling (forward reductions) --- +// +// MLX exposes no window_sum/max/min/product primitive; each is composed +// as pad -> as_strided (sliding-window view) -> reduce over the kernel +// axes. These cores back both the eager NIFs (c_src/ops/pooling.cpp) and +// the compiled program replay, so the two paths can't drift. + +// Contiguous element-strides for a shape, e.g. {B, H, W, C} -> +// {H*W*C, W*C, C, 1}. +inline mx::Strides contiguous_strides(const mx::Shape &shape) { + int rank = static_cast(shape.size()); + mx::Strides out(rank, 1); + for (int i = rank - 2; i >= 0; --i) { + out[i] = out[i + 1] * static_cast(shape[i + 1]); + } + return out; +} + +// Pad `a` with `pad_value` using per-axis lo/hi pads. Returns `a` +// unchanged if all pads are zero (the common path — avoids a copy). +inline mx::array do_pad( + const mx::array &a, + const std::vector &pad_lo, + const std::vector &pad_hi, + const mx::array &pad_value, + mx::Stream &s) { + int rank = static_cast(a.ndim()); + if (pad_lo.size() != static_cast(rank) || + pad_hi.size() != static_cast(rank)) { + throw std::invalid_argument( + "pad: pad_lo/pad_hi length must equal tensor rank " + + std::to_string(rank)); + } + bool any_pad = false; + for (int i = 0; i < rank; ++i) { + if (pad_lo[i] > 0 || pad_hi[i] > 0) { + any_pad = true; + break; + } + } + if (!any_pad) { + return a; + } + + std::vector axes(rank); + std::iota(axes.begin(), axes.end(), 0); + + mx::Shape lo, hi; + lo.reserve(rank); + hi.reserve(rank); + for (int i = 0; i < rank; ++i) { + lo.push_back(static_cast(pad_lo[i])); + hi.push_back(static_cast(pad_hi[i])); + } + + return mx::pad(a, axes, lo, hi, pad_value, "constant", s); +} + +// Build an `as_strided` view with shape `[out_dims..., window_shape...]`. +// `out_dims` is filled with the per-axis output size. +// +// eff_window = (window_shape[i] - 1) * dilations[i] + 1 +// out[i] = (padded_shape[i] - eff_window) / strides[i] + 1 +inline mx::array sliding_windows_view( + const mx::array &padded, + const std::vector &window_shape, + const std::vector &strides, + const std::vector &dilations, + std::vector &out_dims, + mx::Stream &s) { + int rank = static_cast(padded.ndim()); + const auto rank_sz = static_cast(rank); + if (window_shape.size() != rank_sz || strides.size() != rank_sz || + dilations.size() != rank_sz) { + throw std::invalid_argument( + "window: window_shape/strides/dilations length must equal tensor " + "rank " + + std::to_string(rank)); + } + for (int i = 0; i < rank; ++i) { + if (window_shape[i] < 1 || strides[i] < 1 || dilations[i] < 1) { + throw std::invalid_argument( + "window: window dimensions, strides, and dilations must all be " + "positive"); + } + } + const auto &padded_shape = padded.shape(); + auto cs = contiguous_strides(padded_shape); + + out_dims.assign(rank, 0); + mx::Shape new_shape; + mx::Strides new_strides; + new_shape.reserve(2 * rank); + new_strides.reserve(2 * rank); + + for (int i = 0; i < rank; ++i) { + int64_t eff = (window_shape[i] - 1) * dilations[i] + 1; + out_dims[i] = (static_cast(padded_shape[i]) - eff) / strides[i] + 1; + new_shape.push_back(static_cast(out_dims[i])); + } + for (int i = 0; i < rank; ++i) { + new_shape.push_back(static_cast(window_shape[i])); + } + for (int i = 0; i < rank; ++i) { + new_strides.push_back(cs[i] * strides[i]); + } + for (int i = 0; i < rank; ++i) { + new_strides.push_back(cs[i] * dilations[i]); + } + + return mx::as_strided(padded, new_shape, new_strides, 0, s); +} + +enum class WindowReduceKind { Sum, Max, Min, Product }; + +// pad -> sliding-window view -> reduce over the kernel axes. `init_value` +// is the dtype identity (0/1/±inf), used both as the pad fill and (for +// max/min) the reduction's boundary identity. +inline mx::array window_reduce_core( + const mx::array &a, + const std::vector &window_shape, + const std::vector &strides, + const std::vector &pad_lo, + const std::vector &pad_hi, + const std::vector &dilations, + const mx::array &init_value, + WindowReduceKind kind, + mx::Stream &s) { + auto padded = do_pad(a, pad_lo, pad_hi, init_value, s); + std::vector out_dims; + auto view = + sliding_windows_view(padded, window_shape, strides, dilations, out_dims, s); + + // Dilated windows give the kernel axes an `as_strided` stride > 1, so the + // view aliases fewer physical elements than its logical size (overlapping + // strides). MLX's reduction then picks a strided fast path + // (GeneralStridedReduce) that assumes a dense, non-overlapping layout and + // reads `product(shape)` contiguous elements — over-running the buffer and + // returning garbage for windows past the first stride positions (issue + // #175). Materialise the view first: the general copy reads element-by- + // element via the real strides (always in-bounds, since the last window's + // last tap is the last real element), yielding a dense buffer the reducer + // can safely fast-path. Only dilated windows need this; the common + // (non-dilated) pooling path keeps its copy-free strided reduce. + bool dilated = false; + for (int64_t d : dilations) { + if (d > 1) { + dilated = true; + break; + } + } + if (dilated) { + view = mx::contiguous(view, /*allow_col_major=*/false, s); + } + + int rank = static_cast(window_shape.size()); + std::vector reduce_axes(rank); + for (int i = 0; i < rank; ++i) + reduce_axes[i] = rank + i; + + switch (kind) { + case WindowReduceKind::Sum: + return mx::sum(view, reduce_axes, /*keepdims=*/false, s); + case WindowReduceKind::Max: + return mx::max(view, reduce_axes, /*keepdims=*/false, s); + case WindowReduceKind::Min: + return mx::min(view, reduce_axes, /*keepdims=*/false, s); + case WindowReduceKind::Product: + return mx::prod(view, reduce_axes, /*keepdims=*/false, s); + } + throw std::invalid_argument("window_reduce_core: unknown reduce kind"); +} + +// Select-and-scatter — the backward of window_max/window_min (Nx rewrites +// grad(window_max) into window_scatter_max). `is_max` picks argmax vs +// argmin. Tie-break: Nx's select_and_scatter uses `>=`/`<=` (LAST +// occurrence); MLX argmax/argmin give FIRST, so we argmax `mask * pos` +// to recover the last winner. Scatter variants take no dilations. +inline mx::array window_scatter_core( + const mx::array &tensor, + const mx::array &source, + const mx::array &init_value, + const std::vector &window_shape, + const std::vector &strides, + const std::vector &pad_lo, + const std::vector &pad_hi, + bool is_max, + mx::Stream &s) { + int rank = static_cast(window_shape.size()); + auto original_shape = tensor.shape(); + + // 1. Pad input with init_value. + auto padded = do_pad(tensor, pad_lo, pad_hi, init_value, s); + auto padded_shape = padded.shape(); + + // 2. Sliding-window view (dilation is implicitly 1 per axis for scatter). + std::vector dilations(rank, 1); + std::vector out_dims; + auto view = + sliding_windows_view(padded, window_shape, strides, dilations, out_dims, s); + + // 3. Flatten the kernel axes so a single reduction spans the window. + int64_t K = 1; + for (int i = 0; i < rank; ++i) + K *= window_shape[i]; + + mx::Shape flat_view_shape; + flat_view_shape.reserve(rank + 1); + for (int i = 0; i < rank; ++i) + flat_view_shape.push_back(static_cast(out_dims[i])); + flat_view_shape.push_back(static_cast(K)); + + auto flat_view = mx::reshape(view, flat_view_shape, s); + int last_axis = rank; + + // 4. Argmax-with-tie-break (mask*pos picks the last-occurrence winner). + auto selector = is_max + ? mx::max(flat_view, last_axis, /*keepdims=*/true, s) + : mx::min(flat_view, last_axis, /*keepdims=*/true, s); + auto mask = mx::equal(flat_view, selector, s); + + auto pos_1d = mx::arange(0.0, static_cast(K), 1.0, mx::int32, s); + mx::Shape pos_shape(rank + 1, 1); + pos_shape[rank] = static_cast(K); + auto pos = mx::reshape(pos_1d, pos_shape, s); + + auto mask_i = mx::astype(mask, mx::int32, s); + auto mask_pos = mx::multiply(mask_i, pos, s); + auto last_arg = mx::argmax(mask_pos, last_axis, /*keepdims=*/false, s); + + // 5. Decompose the flat kernel index into per-axis kernel indices. + std::vector k_idx; + k_idx.reserve(rank); + for (int i = 0; i < rank; ++i) + k_idx.push_back(last_arg); // placeholder; overwritten below + + mx::array remaining = last_arg; + for (int i = rank - 1; i >= 0; --i) { + auto w_i = mx::array(static_cast(window_shape[i]), mx::int32); + k_idx[i] = mx::remainder(remaining, w_i, s); + if (i > 0) { + remaining = mx::floor_divide(remaining, w_i, s); + } + } + + // 6. Per-axis absolute indices into the padded tensor. + mx::Shape out_shape_s; + out_shape_s.reserve(rank); + for (int i = 0; i < rank; ++i) + out_shape_s.push_back(static_cast(out_dims[i])); + + std::vector abs_indices; + abs_indices.reserve(rank); + for (int i = 0; i < rank; ++i) { + auto base_i = + mx::arange(0.0, static_cast(out_dims[i]), 1.0, mx::int32, s); + mx::Shape bcast(rank, 1); + bcast[i] = static_cast(out_dims[i]); + base_i = mx::reshape(base_i, bcast, s); + auto stride_i = mx::array(static_cast(strides[i]), mx::int32); + auto base_times = mx::multiply(base_i, stride_i, s); + auto bt = mx::broadcast_to(base_times, out_shape_s, s); + abs_indices.push_back(mx::add(bt, k_idx[i], s)); + } + + // 7. Reshape source so each index tuple is a single-point write. + mx::Shape source_reshape; + source_reshape.reserve(2 * rank); + for (int i = 0; i < rank; ++i) + source_reshape.push_back(static_cast(out_dims[i])); + for (int i = 0; i < rank; ++i) + source_reshape.push_back(1); + auto source_r = mx::reshape(source, source_reshape, s); + source_r = mx::astype(source_r, tensor.dtype(), s); + + // 8. Output buffer starts filled with init_value (unselected positions + // retain it; selected positions receive init_value + sum(source)). + auto padded_out = mx::full(padded_shape, init_value, tensor.dtype(), s); + + // 9. Scatter-add all selected contributions in one dispatch. + std::vector axes(rank); + std::iota(axes.begin(), axes.end(), 0); + auto scattered = mx::scatter_add(padded_out, abs_indices, source_r, axes, s); + + // 10. Slice back to the original (unpadded) shape. + mx::Shape slice_start, slice_stop, slice_strides_v; + slice_start.reserve(rank); + slice_stop.reserve(rank); + slice_strides_v.reserve(rank); + for (int i = 0; i < rank; ++i) { + slice_start.push_back(static_cast(pad_lo[i])); + slice_stop.push_back( + static_cast(pad_lo[i] + original_shape[i])); + slice_strides_v.push_back(1); + } + return mx::slice(scattered, slice_start, slice_stop, slice_strides_v, s); +} + } // namespace emily::ops diff --git a/c_src/emily/opcodes.hpp b/c_src/emily/opcodes.hpp index a4e758b..53c7a38 100644 --- a/c_src/emily/opcodes.hpp +++ b/c_src/emily/opcodes.hpp @@ -153,9 +153,24 @@ enum class Opcode : int64_t { Gather = 79, // Stack tensors along a new axis. operands [t0, t1, ...]; iattrs [[axis]] Stack = 80, + // Gather along one axis with a same-rank s32 index tensor. + // operands [input, indices]; iattrs [[axis]] + TakeAlongAxis = 81, + // Window (pooling) reductions: pad -> sliding-window view -> reduce. + // operands [input, init_scalar]; iattrs + // [[window...],[strides...],[pad_lo...],[pad_hi...],[dilations...]] + WindowSum = 82, + WindowMax = 83, + WindowMin = 84, + WindowProduct = 85, + // Window select-and-scatter (MaxPool/MinPool backward). + // operands [input, source, init_scalar]; iattrs + // [[window...],[strides...],[pad_lo...],[pad_hi...]] (no dilations) + WindowScatterMax = 86, + WindowScatterMin = 87, }; -inline constexpr int64_t kOpcodeCount = 81; +inline constexpr int64_t kOpcodeCount = 88; // Quant mode code (Emily.IR @quant_modes) -> MLX mode string. inline std::string qmode_from_code(int64_t code) { @@ -503,6 +518,47 @@ inline mx::array dispatch_op(Opcode op, const std::vector &in, int axis = emily::checked_int(scalar_at(iattrs, 0, "take"), "axis"); return mx::take(in[0], in[1], axis, s); } + case Opcode::TakeAlongAxis: { + if (in.size() != 2) { + throw std::invalid_argument("take_along_axis expects 2 operands, got " + + std::to_string(in.size())); + } + int axis = + emily::checked_int(scalar_at(iattrs, 0, "take_along_axis"), "axis"); + return mx::take_along_axis(in[0], in[1], axis, s); + } + case Opcode::WindowSum: + case Opcode::WindowMax: + case Opcode::WindowMin: + case Opcode::WindowProduct: { + if (in.size() != 2) { + throw std::invalid_argument( + "window reduce expects 2 operands (input, init), got " + + std::to_string(in.size())); + } + auto kind = op == Opcode::WindowSum ? emily::ops::WindowReduceKind::Sum + : op == Opcode::WindowMax ? emily::ops::WindowReduceKind::Max + : op == Opcode::WindowMin ? emily::ops::WindowReduceKind::Min + : emily::ops::WindowReduceKind::Product; + return emily::ops::window_reduce_core( + in[0], attr_at(iattrs, 0, "window"), attr_at(iattrs, 1, "window"), + attr_at(iattrs, 2, "window"), attr_at(iattrs, 3, "window"), + attr_at(iattrs, 4, "window"), in[1], kind, s); + } + case Opcode::WindowScatterMax: + case Opcode::WindowScatterMin: { + if (in.size() != 3) { + throw std::invalid_argument( + "window scatter expects 3 operands (input, source, init), got " + + std::to_string(in.size())); + } + bool is_max = op == Opcode::WindowScatterMax; + return emily::ops::window_scatter_core( + in[0], in[1], in[2], attr_at(iattrs, 0, "window_scatter"), + attr_at(iattrs, 1, "window_scatter"), + attr_at(iattrs, 2, "window_scatter"), + attr_at(iattrs, 3, "window_scatter"), is_max, s); + } case Opcode::Concatenate: { if (in.empty()) { throw std::invalid_argument("concatenate expects >= 1 operand"); diff --git a/c_src/ops/pooling.cpp b/c_src/ops/pooling.cpp index 8615f57..bbd65c5 100644 --- a/c_src/ops/pooling.cpp +++ b/c_src/ops/pooling.cpp @@ -15,6 +15,7 @@ // back to the unpadded shape. #include "../emily/async.hpp" +#include "../emily/op_cores.hpp" #include "../emily/tensor.hpp" #include "../emily/worker.hpp" @@ -36,130 +37,20 @@ using emily::WorkerThread; namespace { -// -------------------- Shared helpers -------------------- - -// Contiguous element-strides for a shape, e.g. {B, H, W, C} -> -// {H*W*C, W*C, C, 1}. -mx::Strides contiguous_strides(const mx::Shape &shape) { - int rank = static_cast(shape.size()); - mx::Strides out(rank, 1); - for (int i = rank - 2; i >= 0; --i) { - out[i] = out[i + 1] * static_cast(shape[i + 1]); - } - return out; -} - -// Pad `a` with `pad_value` using per-axis lo/hi pads. Returns `a` -// unchanged if all pads are zero (the common path — avoids a pointless -// copy). -mx::array do_pad( - const mx::array &a, - const std::vector &pad_lo, - const std::vector &pad_hi, - const mx::array &pad_value, - mx::Stream &s) { - int rank = static_cast(a.ndim()); - // A direct Native call can pass pad vectors shorter than the tensor - // rank; the per-axis loops below would then read out of bounds. - if (pad_lo.size() != static_cast(rank) || - pad_hi.size() != static_cast(rank)) { - throw std::invalid_argument( - "pad: pad_lo/pad_hi length must equal tensor rank " + - std::to_string(rank)); - } - bool any_pad = false; - for (int i = 0; i < rank; ++i) { - if (pad_lo[i] > 0 || pad_hi[i] > 0) { - any_pad = true; - break; - } - } - if (!any_pad) { - return a; - } - - std::vector axes(rank); - std::iota(axes.begin(), axes.end(), 0); - - mx::Shape lo, hi; - lo.reserve(rank); - hi.reserve(rank); - for (int i = 0; i < rank; ++i) { - lo.push_back(static_cast(pad_lo[i])); - hi.push_back(static_cast(pad_hi[i])); - } - - return mx::pad(a, axes, lo, hi, pad_value, "constant", s); -} - -// Build an `as_strided` view with shape `[out_dims..., window_shape...]`. -// Output `out_dims` is filled with the per-axis output size. -// -// Output shape formula (per axis): -// eff_window = (window_shape[i] - 1) * dilations[i] + 1 -// out[i] = (padded_shape[i] - eff_window) / strides[i] + 1 -// -// Strides (in elements, relative to the padded tensor's contiguous -// layout — `as_strided` forces its input to be contiguous internally): -// out-axis i: contiguous_stride[i] * strides[i] -// kernel-ax i: contiguous_stride[i] * dilations[i] -mx::array sliding_windows_view( - const mx::array &padded, - const std::vector &window_shape, - const std::vector &strides, - const std::vector &dilations, - std::vector &out_dims, - mx::Stream &s) { - int rank = static_cast(padded.ndim()); - // A direct Native call can pass window/stride/dilation vectors that - // don't match the tensor rank (out-of-bounds indexing below) or a zero - // stride (the `/ strides[i]` out-shape divide is an integer SIGFPE that - // bypasses the async catch ladder and crashes the BEAM). - const auto rank_sz = static_cast(rank); - if (window_shape.size() != rank_sz || strides.size() != rank_sz || - dilations.size() != rank_sz) { - throw std::invalid_argument( - "window: window_shape/strides/dilations length must equal tensor " - "rank " + - std::to_string(rank)); - } - for (int i = 0; i < rank; ++i) { - if (window_shape[i] < 1 || strides[i] < 1 || dilations[i] < 1) { - throw std::invalid_argument( - "window: window dimensions, strides, and dilations must all be " - "positive"); - } - } - const auto &padded_shape = padded.shape(); - auto cs = contiguous_strides(padded_shape); - - out_dims.assign(rank, 0); - mx::Shape new_shape; - mx::Strides new_strides; - new_shape.reserve(2 * rank); - new_strides.reserve(2 * rank); - - for (int i = 0; i < rank; ++i) { - int64_t eff = (window_shape[i] - 1) * dilations[i] + 1; - out_dims[i] = (static_cast(padded_shape[i]) - eff) / strides[i] + 1; - new_shape.push_back(static_cast(out_dims[i])); - } - for (int i = 0; i < rank; ++i) { - new_shape.push_back(static_cast(window_shape[i])); - } - for (int i = 0; i < rank; ++i) { - new_strides.push_back(cs[i] * strides[i]); - } - for (int i = 0; i < rank; ++i) { - new_strides.push_back(cs[i] * dilations[i]); - } - - return mx::as_strided(padded, new_shape, new_strides, 0, s); -} +// The window cores (do_pad / sliding_windows_view / window_reduce_core) +// live in emily/op_cores.hpp so the eager NIFs below and the Expr-compiler +// program replay (c_src/program.cpp) share one implementation and can't +// numerically drift. window_scatter_impl below still composes do_pad + +// sliding_windows_view from there. +using emily::ops::do_pad; +using emily::ops::sliding_windows_view; +using emily::ops::window_reduce_core; +using emily::ops::window_scatter_core; +using emily::ops::WindowReduceKind; // -------------------- Reductions -------------------- -#define EMILY_WINDOW_REDUCE(op_name, mlx_fn) \ +#define EMILY_WINDOW_REDUCE(op_name, kind) \ fine::Term op_name##_nif( \ ErlNifEnv *env, \ fine::ResourcePtr w, \ @@ -175,165 +66,25 @@ mx::array sliding_windows_view( strides = std::move(strides), pad_lo = std::move(pad_lo), \ pad_hi = std::move(pad_hi), dilations = std::move(dilations), \ init_value = std::move(init_value)](mx::Stream &s) { \ - auto padded = do_pad(t->array, pad_lo, pad_hi, init_value->array, s);\ - std::vector out_dims; \ - auto view = sliding_windows_view(padded, window_shape, strides, \ - dilations, out_dims, s); \ - int rank = static_cast(window_shape.size()); \ - std::vector reduce_axes(rank); \ - for (int i = 0; i < rank; ++i) \ - reduce_axes[i] = rank + i; \ - return wrap(mlx_fn(view, reduce_axes, /*keepdims=*/false, s)); \ + return wrap(window_reduce_core( \ + t->array, window_shape, strides, pad_lo, pad_hi, dilations, \ + init_value->array, kind, s)); \ }); \ } \ FINE_NIF(op_name##_nif, 0); -EMILY_WINDOW_REDUCE(window_sum, mx::sum) -EMILY_WINDOW_REDUCE(window_max, mx::max) -EMILY_WINDOW_REDUCE(window_min, mx::min) -EMILY_WINDOW_REDUCE(window_product, mx::prod) +EMILY_WINDOW_REDUCE(window_sum, WindowReduceKind::Sum) +EMILY_WINDOW_REDUCE(window_max, WindowReduceKind::Max) +EMILY_WINDOW_REDUCE(window_min, WindowReduceKind::Min) +EMILY_WINDOW_REDUCE(window_product, WindowReduceKind::Product) #undef EMILY_WINDOW_REDUCE // -------------------- Scatter variants -------------------- - -// Shared body: select-and-scatter. `is_max` picks between argmax -// (for window_scatter_max) and argmin (for window_scatter_min). // -// Tie-break semantics: Nx uses `>=` / `<=` in its select_and_scatter, -// i.e. the LAST-occurrence winner. MLX's argmax/argmin return -// FIRST-occurrence. We build `mask_pos = (flat_view == selector) * -// arange(K)` and argmax that — for tied positions the pos multiplier -// makes the later index strictly larger, giving last-occurrence. -mx::array window_scatter_impl( - const mx::array &tensor, - const mx::array &source, - const mx::array &init_value, - const std::vector &window_shape, - const std::vector &strides, - const std::vector &pad_lo, - const std::vector &pad_hi, - bool is_max, - mx::Stream &s) { - int rank = static_cast(window_shape.size()); - auto original_shape = tensor.shape(); - - // 1. Pad input with init_value. - auto padded = do_pad(tensor, pad_lo, pad_hi, init_value, s); - auto padded_shape = padded.shape(); - - // 2. Sliding-window view. Scatter variants don't take dilations in - // Nx's API, so dilation is implicitly 1 per axis. - std::vector dilations(rank, 1); - std::vector out_dims; - auto view = - sliding_windows_view(padded, window_shape, strides, dilations, out_dims, s); - - // 3. Flatten the kernel axes so we can argmax across the whole window - // in a single reduction. - int64_t K = 1; - for (int i = 0; i < rank; ++i) - K *= window_shape[i]; - - mx::Shape flat_view_shape; - flat_view_shape.reserve(rank + 1); - for (int i = 0; i < rank; ++i) - flat_view_shape.push_back(static_cast(out_dims[i])); - flat_view_shape.push_back(static_cast(K)); - - auto flat_view = mx::reshape(view, flat_view_shape, s); - int last_axis = rank; - - // 4. Argmax-with-tie-break. `selector` is the per-window max/min; mask - // is 1 where the kernel element equals the winner; mask*pos gives - // later-matching positions a higher value, so argmax picks the - // last-occurrence kernel index. - auto selector = is_max - ? mx::max(flat_view, last_axis, /*keepdims=*/true, s) - : mx::min(flat_view, last_axis, /*keepdims=*/true, s); - auto mask = mx::equal(flat_view, selector, s); - - auto pos_1d = mx::arange(0.0, static_cast(K), 1.0, mx::int32, s); - mx::Shape pos_shape(rank + 1, 1); - pos_shape[rank] = static_cast(K); - auto pos = mx::reshape(pos_1d, pos_shape, s); - - auto mask_i = mx::astype(mask, mx::int32, s); - auto mask_pos = mx::multiply(mask_i, pos, s); - auto last_arg = mx::argmax(mask_pos, last_axis, /*keepdims=*/false, s); - - // 5. Decompose flat kernel index into per-axis kernel indices. - // k_idx[R-1] = last_arg % window[R-1]; - // k_idx[R-2] = (last_arg / window[R-1]) % window[R-2]; ... - std::vector k_idx; - k_idx.reserve(rank); - for (int i = 0; i < rank; ++i) - k_idx.push_back(last_arg); // placeholder; overwritten below - - mx::array remaining = last_arg; - for (int i = rank - 1; i >= 0; --i) { - auto w_i = mx::array(static_cast(window_shape[i]), mx::int32); - k_idx[i] = mx::remainder(remaining, w_i, s); - if (i > 0) { - remaining = mx::floor_divide(remaining, w_i, s); - } - } - - // 6. Per-axis absolute indices into the padded tensor: - // abs_idx[i][out_coord] = out_coord[i] * stride[i] + k_idx[i][out_coord] - mx::Shape out_shape_s; - out_shape_s.reserve(rank); - for (int i = 0; i < rank; ++i) - out_shape_s.push_back(static_cast(out_dims[i])); - - std::vector abs_indices; - abs_indices.reserve(rank); - for (int i = 0; i < rank; ++i) { - auto base_i = - mx::arange(0.0, static_cast(out_dims[i]), 1.0, mx::int32, s); - mx::Shape bcast(rank, 1); - bcast[i] = static_cast(out_dims[i]); - base_i = mx::reshape(base_i, bcast, s); - auto stride_i = mx::array(static_cast(strides[i]), mx::int32); - auto base_times = mx::multiply(base_i, stride_i, s); - auto bt = mx::broadcast_to(base_times, out_shape_s, s); - abs_indices.push_back(mx::add(bt, k_idx[i], s)); - } - - // 7. Reshape source to IDX_SHAPE + [1]*rank so MLX scatter_add treats - // each index tuple as a single-point write. - mx::Shape source_reshape; - source_reshape.reserve(2 * rank); - for (int i = 0; i < rank; ++i) - source_reshape.push_back(static_cast(out_dims[i])); - for (int i = 0; i < rank; ++i) - source_reshape.push_back(1); - auto source_r = mx::reshape(source, source_reshape, s); - source_r = mx::astype(source_r, tensor.dtype(), s); - - // 8. Output buffer starts filled with init_value. Matches Nx's - // select_and_scatter: unselected positions retain init_value; - // selected positions receive init_value + sum(source values). - auto padded_out = mx::full(padded_shape, init_value, tensor.dtype(), s); - - // 9. Scatter-add all selected contributions in one dispatch. - std::vector axes(rank); - std::iota(axes.begin(), axes.end(), 0); - auto scattered = mx::scatter_add(padded_out, abs_indices, source_r, axes, s); - - // 10. Slice back to original (unpadded) shape. - mx::Shape slice_start, slice_stop, slice_strides_v; - slice_start.reserve(rank); - slice_stop.reserve(rank); - slice_strides_v.reserve(rank); - for (int i = 0; i < rank; ++i) { - slice_start.push_back(static_cast(pad_lo[i])); - slice_stop.push_back( - static_cast(pad_lo[i] + original_shape[i])); - slice_strides_v.push_back(1); - } - return mx::slice(scattered, slice_start, slice_stop, slice_strides_v, s); -} +// window_scatter_core (the MaxPool/MinPool backward) lives in +// emily/op_cores.hpp so the eager NIFs below and the Expr-compiler program +// replay share one implementation. fine::Term window_scatter_max_nif( ErlNifEnv *env, @@ -351,7 +102,7 @@ fine::Term window_scatter_max_nif( window_shape = std::move(window_shape), strides = std::move(strides), pad_lo = std::move(pad_lo), pad_hi = std::move(pad_hi)](mx::Stream &s) { - return wrap(window_scatter_impl( + return wrap(window_scatter_core( t->array, source->array, init_value->array, window_shape, strides, pad_lo, pad_hi, /*is_max=*/true, s)); }); @@ -374,7 +125,7 @@ fine::Term window_scatter_min_nif( window_shape = std::move(window_shape), strides = std::move(strides), pad_lo = std::move(pad_lo), pad_hi = std::move(pad_hi)](mx::Stream &s) { - return wrap(window_scatter_impl( + return wrap(window_scatter_core( t->array, source->array, init_value->array, window_shape, strides, pad_lo, pad_hi, /*is_max=*/false, s)); }); diff --git a/lib/emily/ir.ex b/lib/emily/ir.ex index 34b5b34..3f3f331 100644 --- a/lib/emily/ir.ex +++ b/lib/emily/ir.ex @@ -132,7 +132,20 @@ defmodule Emily.IR do cummin: 78, # multi-axis gather: operands [input, idx0, ...]; iattrs [[axes],[slice_sizes]] gather: 79, - stack: 80 + stack: 80, + # take_along_axis: gather along one axis with a same-rank s32 index + # tensor. operands [input, indices]; iattrs [[axis]]. + take_along_axis: 81, + # window (pooling) reductions. operands [input, init_scalar]; iattrs + # [[window],[strides],[pad_lo],[pad_hi],[dilations]]. + window_sum: 82, + window_max: 83, + window_min: 84, + window_product: 85, + # window select-and-scatter (pooling backward). operands + # [input, source, init]; iattrs [[window],[strides],[pad_lo],[pad_hi]]. + window_scatter_max: 86, + window_scatter_min: 87 } # Quant mode string -> code; decoded by qmode_from_code in @@ -163,6 +176,14 @@ defmodule Emily.IR do outputs: [ref()] } + @doc """ + The opcode name -> wire-value map. Exposed for the opcode-parity test, + which checks these stay in lockstep with the `Opcode` enum and + `kOpcodeCount` in `c_src/emily/opcodes.hpp`. + """ + @spec opcodes() :: %{atom() => non_neg_integer()} + def opcodes, do: @opcodes + @doc "Numeric wire value for an opcode name." @spec opcode(atom()) :: non_neg_integer() def opcode(name) when is_map_key(@opcodes, name), do: Map.fetch!(@opcodes, name) @@ -727,6 +748,57 @@ defmodule Emily.IR do coerce(r, type, state) end + # Window (pooling) reductions (window_sum/max/min/product). Mirrors + # Emily.Backend.apply_window_reduce/6: pad with the dtype identity, build + # the sliding-window view, reduce over the kernel axes. The init scalar + # is baked as a const operand so float ±inf and integer min/max are all + # handled exactly as the eager `identity_ref`. Operands [input, init]; + # iattrs [[window],[strides],[pad_lo],[pad_hi],[dilations]]. + defp lower_op(%T{data: %Nx.Defn.Expr{op: op, args: [a, window_dimensions, opts]}} = t, state) + when op in [:window_sum, :window_max, :window_min, :window_product] do + {ra, state} = lower_node(a, state) + {init_ref, state} = window_identity(op, a.type, state) + + rank = tuple_size(a.shape) + window = Tuple.to_list(window_dimensions) + strides = window_per_axis(opts[:strides], rank, 1) + dilations = window_per_axis(opts[:window_dilations], rank, 1) + {pad_lo, pad_hi} = window_padding(opts[:padding], rank) + + emit_coerced(state, op, [ra, init_ref], [window, strides, pad_lo, pad_hi, dilations], t.type) + end + + # Window select-and-scatter (window_scatter_max/min) — the MaxPool/MinPool + # backward. Mirrors Emily.Backend.apply_window_scatter/7: operands + # [input, source, init]; iattrs [[window],[strides],[pad_lo],[pad_hi]] + # (no dilations). The init scalar comes from the Expr (Nx's grad rule), + # coerced to the output dtype. + defp lower_op( + %T{data: %Nx.Defn.Expr{op: op, args: [tin, source, init, window_dimensions, opts]}} = t, + state + ) + when op in [:window_scatter_max, :window_scatter_min] do + {rt, state} = lower_node(tin, state) + {rs, state} = lower_node(source, state) + {ri, state} = lower_node(init, state) + {ri, state} = emit(state, :astype, [ri], [[dtype_code(t.type)]]) + + rank = tuple_size(tin.shape) + window = Tuple.to_list(window_dimensions) + strides = window_per_axis(opts[:strides], rank, 1) + {pad_lo, pad_hi} = window_padding(opts[:padding], rank) + + emit_coerced(state, op, [rt, rs, ri], [window, strides, pad_lo, pad_hi], t.type) + end + + # Nx.reverse along one or more axes (the conv backward flips the kernel). + # Reversing is order-independent across axes, so chain a single-axis + # `flip` (mx negative-stride slice) per axis. Empty axes => identity. + defp lower_op(%T{data: %Nx.Defn.Expr{op: :reverse, args: [a, axes]}}, state) do + {ra, state} = lower_node(a, state) + Enum.reduce(axes, {ra, state}, fn axis, {r, st} -> emit(st, :flip, [r], [[axis]]) end) + end + # cond: raw args [clauses, last], clauses = [{pred, body}, ...]. Lower to # a select chain `where(p1, b1, where(p2, b2, ... last))`. ALL branches # are evaluated (Nx branches are side-effect-free and shape-compatible); @@ -915,6 +987,16 @@ defmodule Emily.IR do emit_coerced(state, :take, [ri, rx], [[axis]], t.type) end + # Nx.take_along_axis (Nx.Block.TakeAlongAxis). Mirrors + # Emily.Backend.native_take_along_axis/4: cast indices to s32, then + # mx::take_along_axis along `axis`. + defp lower_block(%Nx.Block.TakeAlongAxis{axis: axis}, [input, indices], _expr, t, state) do + {ri, state} = lower_node(input, state) + {rx, state} = lower_node(indices, state) + {rx, state} = emit(state, :astype, [rx], [[dtype_code({:s, 32})]]) + emit_coerced(state, :take_along_axis, [ri, rx], [[axis]], t.type) + end + # Cumulative families. Like Emily.Backend.block/4, the last-axis case uses # the native MLX `cumsum`/`cumprod`/`cummax`/`cummin` kernel; interior axes # (which MLX can't always factor) fall back to the block's composed @@ -1028,6 +1110,49 @@ defmodule Emily.IR do # holds materialized literal constants (and iota); `:capture` holds # embedded `:tensor` weights. Both go through from_binary once at # lower time and are never re-shipped. + # Bake the dtype identity for a window reduce as a scalar const operand + # (mirrors Emily.Backend.identity_ref/2): 0 for sum, 1 for product, + # ±inf for float max/min, dtype min/max for integer max/min. + defp window_identity(op, type, state) do + materialize_const(window_identity_scalar(op, type), {}, type, state) + end + + defp window_identity_scalar(:window_sum, type), + do: Nx.tensor(0, type: type, backend: Nx.BinaryBackend) + + defp window_identity_scalar(:window_product, type), + do: Nx.tensor(1, type: type, backend: Nx.BinaryBackend) + + defp window_identity_scalar(:window_max, {kind, _} = type) when kind in [:f, :bf], + do: Nx.tensor(:neg_infinity, type: type, backend: Nx.BinaryBackend) + + defp window_identity_scalar(:window_min, {kind, _} = type) when kind in [:f, :bf], + do: Nx.tensor(:infinity, type: type, backend: Nx.BinaryBackend) + + defp window_identity_scalar(:window_max, {kind, bits} = type) when kind in [:s, :u] do + value = if kind == :u, do: 0, else: -Bitwise.bsl(1, bits - 1) + Nx.tensor(value, type: type, backend: Nx.BinaryBackend) + end + + defp window_identity_scalar(:window_min, {kind, bits} = type) when kind in [:s, :u] do + value = if kind == :u, do: Bitwise.bsl(1, bits) - 1, else: Bitwise.bsl(1, bits - 1) - 1 + Nx.tensor(value, type: type, backend: Nx.BinaryBackend) + end + + # Per-axis strides / dilations (mirror Emily.Backend.normalize_per_axis/3). + defp window_per_axis(nil, rank, default), do: List.duplicate(default, rank) + defp window_per_axis(n, rank, _default) when is_integer(n), do: List.duplicate(n, rank) + defp window_per_axis(list, _rank, _default) when is_list(list), do: list + + # Split resolved `[{lo, hi}, ...]` padding into two lists (mirror + # Emily.Backend.split_padding/2). + defp window_padding(pairs, _rank) when is_list(pairs) do + pairs |> Enum.map(fn {lo, hi} -> {lo, hi} end) |> Enum.unzip() + end + + defp window_padding(_other, rank), + do: {List.duplicate(0, rank), List.duplicate(0, rank)} + defp materialize_const(tensor, shape, type, state) do ref = Emily.Native.from_binary(Nx.to_binary(tensor), Tuple.to_list(shape), type) idx = state.n_consts diff --git a/notebooks/distilbert_qa.livemd b/livebooks/distilbert_qa.livemd similarity index 100% rename from notebooks/distilbert_qa.livemd rename to livebooks/distilbert_qa.livemd diff --git a/notebooks/fast_kernels.livemd b/livebooks/fast_kernels.livemd similarity index 100% rename from notebooks/fast_kernels.livemd rename to livebooks/fast_kernels.livemd diff --git a/notebooks/mnist_training.livemd b/livebooks/mnist_training.livemd similarity index 98% rename from notebooks/mnist_training.livemd rename to livebooks/mnist_training.livemd index 5500b79..52311b2 100644 --- a/notebooks/mnist_training.livemd +++ b/livebooks/mnist_training.livemd @@ -17,7 +17,7 @@ Mix.install( ## Overview -The other notebooks in this repo run pretrained models forward. This +The other livebooks in this repo run pretrained models forward. This one trains from scratch — a small Axon MLP on MNIST, lowered through `Emily.Compiler` so every step of the backward pass dispatches to MLX. It's the end-to-end exercise of the `Nx.Defn.grad` chain that diff --git a/notebooks/modernbert_classification.livemd b/livebooks/modernbert_classification.livemd similarity index 100% rename from notebooks/modernbert_classification.livemd rename to livebooks/modernbert_classification.livemd diff --git a/notebooks/nomic_embeddings.livemd b/livebooks/nomic_embeddings.livemd similarity index 100% rename from notebooks/nomic_embeddings.livemd rename to livebooks/nomic_embeddings.livemd diff --git a/notebooks/qwen3_quantized.livemd b/livebooks/qwen3_quantized.livemd similarity index 100% rename from notebooks/qwen3_quantized.livemd rename to livebooks/qwen3_quantized.livemd diff --git a/notebooks/smollm3_chat.livemd b/livebooks/smollm3_chat.livemd similarity index 100% rename from notebooks/smollm3_chat.livemd rename to livebooks/smollm3_chat.livemd diff --git a/notebooks/whisper_transcription.livemd b/livebooks/whisper_transcription.livemd similarity index 100% rename from notebooks/whisper_transcription.livemd rename to livebooks/whisper_transcription.livemd diff --git a/mix.exs b/mix.exs index fa0d800..31c4515 100644 --- a/mix.exs +++ b/mix.exs @@ -230,14 +230,14 @@ defmodule Emily.MixProject do "ARCHITECTURE.md", "ROADMAP.md", "CHANGELOG.md", - "notebooks/distilbert_qa.livemd", - "notebooks/qwen3_quantized.livemd", - "notebooks/nomic_embeddings.livemd", - "notebooks/smollm3_chat.livemd", - "notebooks/modernbert_classification.livemd", - "notebooks/mnist_training.livemd", - "notebooks/whisper_transcription.livemd", - "notebooks/fast_kernels.livemd" + "livebooks/distilbert_qa.livemd", + "livebooks/qwen3_quantized.livemd", + "livebooks/nomic_embeddings.livemd", + "livebooks/smollm3_chat.livemd", + "livebooks/modernbert_classification.livemd", + "livebooks/mnist_training.livemd", + "livebooks/whisper_transcription.livemd", + "livebooks/fast_kernels.livemd" ], groups_for_extras: [ README: ~r{README.md}, @@ -246,7 +246,7 @@ defmodule Emily.MixProject do "ROADMAP.md", "CHANGELOG.md" ], - Notebooks: ~r{^notebooks/} + Livebooks: ~r{^livebooks/} ], groups_for_modules: [ Core: [Emily, Emily.Backend, Emily.Compiler], diff --git a/scripts/test-livebooks.sh b/scripts/test-livebooks.sh index 74655a0..5bb55ed 100755 --- a/scripts/test-livebooks.sh +++ b/scripts/test-livebooks.sh @@ -1,7 +1,7 @@ #!/usr/bin/env bash # Smoke-test the example livebooks against the LOCAL emily checkout. # -# For each notebook in notebooks/, this extracts its Elixir cells, +# For each notebook in livebooks/, this extracts its Elixir cells, # repoints the `{:emily, "~> x"}` Mix.install dependency at this repo (as # a `path:` dep, so the notebook exercises the working tree — including a # from-source NIF build), and runs the result headlessly with `elixir`. @@ -28,7 +28,7 @@ set -uo pipefail REPO="$(cd "$(dirname "$0")/.." && pwd)" -NB_DIR="$REPO/notebooks" +NB_DIR="$REPO/livebooks" TIMEOUT="${LIVEBOOK_TIMEOUT:-1200}" SKIP=" ${LIVEBOOK_SKIP:-} " diff --git a/test/emily/backend_window_test.exs b/test/emily/backend_window_test.exs index 9f3ba31..f2f4025 100644 --- a/test/emily/backend_window_test.exs +++ b/test/emily/backend_window_test.exs @@ -160,4 +160,36 @@ defmodule Emily.BackendWindowTest do assert_close(result, bin(t), tol: 1.0e-5) end end + + # Regression for issue #175. 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). On small inputs MLX's + # strided-reduce fast path read past the aliased buffer and returned + # garbage for windows past the first stride positions. These shapes are + # deliberately tiny so the over-read crosses the allocation — the larger + # dilated cases above happened to land on valid data and masked the bug. + describe "dilated windows over small tensors (issue #175 regression)" do + test "1-D kernel, dilation 2" do + t = fixt({8}) + run(:window_sum, t, {3}, window_dilations: [2]) + run(:window_max, t, {3}, window_dilations: [2]) + run(:window_min, t, {3}, window_dilations: [2]) + end + + test "2-D row vector, dilation on the inner axis" do + t = fixt({1, 8}) + run(:window_sum, t, {1, 3}, window_dilations: [1, 2]) + run(:window_max, t, {1, 3}, window_dilations: [1, 2]) + run(:window_min, t, {1, 3}, window_dilations: [1, 2]) + end + + test "window_product, dilation 2" do + t = + Nx.iota({6}, type: {:f, 32}, backend: Nx.BinaryBackend) + |> Nx.multiply(0.2) + |> Nx.add(1.0) + + run(:window_product, t, {3}, window_dilations: [2]) + end + end end diff --git a/test/emily/compiler_equivalence_test.exs b/test/emily/compiler_equivalence_test.exs index 45a7467..987bdbb 100644 --- a/test/emily/compiler_equivalence_test.exs +++ b/test/emily/compiler_equivalence_test.exs @@ -140,6 +140,13 @@ defmodule Emily.CompilerEquivalenceTest do x = et([1.0, 2.0, 3.0]) assert_equiv(fn t -> Nx.broadcast(t, {2, 3}) end, [x]) end + + test "reverse along one and multiple axes matches" do + x = et([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + assert_equiv(fn t -> Nx.reverse(t, axes: [1]) end, [x]) + assert_equiv(fn t -> Nx.reverse(t, axes: [0, 1]) end, [x]) + assert_equiv(fn t -> Nx.reverse(t) end, [x]) + end end describe "dot / matmul" do @@ -368,6 +375,107 @@ defmodule Emily.CompilerEquivalenceTest do end end + describe "take_along_axis" do + test "along the last axis matches the Evaluator" do + x = et([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]]) + idx = Nx.tensor([[3, 0], [1, 2], [0, 3]], type: :s64, backend: Emily.Backend) + assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 1) end, [x, idx]) + end + + test "along axis 0 matches" do + x = et([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + idx = Nx.tensor([[0, 1, 0], [1, 0, 1], [0, 0, 1]], type: :s64, backend: Emily.Backend) + assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 0) end, [x, idx]) + end + + test "3-D gather along the last axis (transformer-shaped) matches" do + x = Nx.iota({1, 2, 4}, type: :f32, backend: Emily.Backend) |> Nx.divide(8.0) + idx = Nx.tensor([[[3, 1], [0, 2]]], type: :s32, backend: Emily.Backend) + assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 2) end, [x, idx]) + end + end + + describe "window reductions (pooling forward)" do + test "2x2 maxpool / sumpool / minpool (CNN-shaped) match the Evaluator" do + # {batch, channels, h, w}; pool only the spatial axes, stride 2. + x = Nx.iota({1, 2, 4, 4}, type: :f32, backend: Emily.Backend) |> Nx.divide(16.0) + + assert_equiv(fn t -> Nx.window_max(t, {1, 1, 2, 2}, strides: [1, 1, 2, 2]) end, [x]) + assert_equiv(fn t -> Nx.window_sum(t, {1, 1, 2, 2}, strides: [1, 1, 2, 2]) end, [x]) + assert_equiv(fn t -> Nx.window_min(t, {1, 1, 2, 2}, strides: [1, 1, 2, 2]) end, [x]) + end + + test "maxpool with padding (boundary identity) matches" do + x = Nx.iota({1, 1, 5, 5}, type: :f32, backend: Emily.Backend) |> Nx.divide(25.0) + + assert_equiv( + fn t -> + Nx.window_max(t, {1, 1, 3, 3}, + strides: [1, 1, 2, 2], + padding: [{0, 0}, {0, 0}, {1, 1}, {1, 1}] + ) + end, + [x] + ) + end + + test "window_product matches" do + x = + Nx.iota({1, 8}, type: :f32, backend: Emily.Backend) |> Nx.divide(8.0) |> Nx.add(1.0) + + assert_equiv(fn t -> Nx.window_product(t, {1, 2}, strides: [1, 1]) end, [x]) + end + + test "dilated windows match (issue #175)" do + # Tiny input + dilation 2 on the inner axis: the sliding-window view + # over-reads its aliased buffer unless materialised before the reduce. + # Pre-fix the native path returned garbage here while the Evaluator + # returned a deterministic 0.0, so the two diverged. + x = + Nx.iota({1, 8}, type: :f32, backend: Emily.Backend) |> Nx.divide(8.0) |> Nx.add(1.0) + + assert_equiv(fn t -> Nx.window_max(t, {1, 3}, window_dilations: [1, 2]) end, [x]) + assert_equiv(fn t -> Nx.window_sum(t, {1, 3}, window_dilations: [1, 2]) end, [x]) + assert_equiv(fn t -> Nx.window_min(t, {1, 3}, window_dilations: [1, 2]) end, [x]) + end + end + + describe "window scatter (pooling backward)" do + # The MaxPool/MinPool backward: scatter the upstream gradient `source` + # (one per pooled window) into the argmax/argmin position of each window. + test "window_scatter_max / window_scatter_min match the Evaluator" do + t = Nx.iota({1, 1, 4, 4}, type: :f32, backend: Emily.Backend) |> Nx.divide(16.0) + source = Nx.iota({1, 1, 2, 2}, type: :f32, backend: Emily.Backend) |> Nx.add(1.0) + + assert_equiv( + fn t, src -> + Nx.window_scatter_max(t, src, 0.0, {1, 1, 2, 2}, strides: [1, 1, 2, 2]) + end, + [t, source] + ) + + assert_equiv( + fn t, src -> + Nx.window_scatter_min(t, src, 0.0, {1, 1, 2, 2}, strides: [1, 1, 2, 2]) + end, + [t, source] + ) + end + + test "grad(window_max) lowers (the maxpool backward path) and matches" do + x = Nx.iota({1, 1, 4, 4}, type: :f32, backend: Emily.Backend) |> Nx.divide(16.0) + + assert_equiv( + fn t -> + Nx.Defn.grad(t, fn t -> + t |> Nx.window_max({1, 1, 2, 2}, strides: [1, 1, 2, 2]) |> Nx.sum() + end) + end, + [x] + ) + end + end + describe "dynamic put_slice (KV-cache write)" do test "put_slice at a runtime offset matches the Evaluator" do # {batch, n_kv_heads, max_len, head_dim} KV buffer; write one token. diff --git a/test/emily/conformance/distilbert_test.exs b/test/emily/conformance/distilbert_test.exs index 4929472..83f050a 100644 --- a/test/emily/conformance/distilbert_test.exs +++ b/test/emily/conformance/distilbert_test.exs @@ -33,7 +33,8 @@ defmodule Emily.Conformance.DistilbertTest do use ExUnit.Case, async: false - import Emily.ConformanceHelper, only: [assert_all_close: 2, assert_all_close: 3] + import Emily.ConformanceHelper, + only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2, mode_test: 3] alias Emily.Bumblebee.FastKernels @@ -52,7 +53,7 @@ defmodule Emily.Conformance.DistilbertTest do :ok end - test ":base" do + mode_test ":base" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertModel"}) @@ -63,7 +64,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.hidden_state) == {1, 10, 32} @@ -75,7 +76,7 @@ defmodule Emily.Conformance.DistilbertTest do ) end - test ":for_masked_language_modeling" do + mode_test ":for_masked_language_modeling" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertForMaskedLM"}) @@ -86,7 +87,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 10, 1124} @@ -98,7 +99,7 @@ defmodule Emily.Conformance.DistilbertTest do ) end - test ":for_sequence_classification" do + mode_test ":for_sequence_classification" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-DistilBertForSequenceClassification"} @@ -111,14 +112,14 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 2} assert_all_close(outputs.logits, Nx.tensor([[-0.0047, -0.0103]])) end - test ":for_token_classification" do + mode_test ":for_token_classification" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-DistilBertForTokenClassification"} @@ -131,7 +132,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 10, 2} @@ -141,7 +142,7 @@ defmodule Emily.Conformance.DistilbertTest do ) end - test ":for_question_answering" do + mode_test ":for_question_answering" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-DistilBertForQuestionAnswering"} @@ -154,7 +155,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.start_logits) == {1, 10} assert Nx.shape(outputs.end_logits) == {1, 10} @@ -170,7 +171,7 @@ defmodule Emily.Conformance.DistilbertTest do ) end - test ":for_multiple_choice" do + mode_test ":for_multiple_choice" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-DistilBertForMultipleChoice"} @@ -183,7 +184,7 @@ defmodule Emily.Conformance.DistilbertTest do "attention_mask" => Nx.tensor([[[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 1} @@ -228,15 +229,22 @@ defmodule Emily.Conformance.DistilbertTest do # (vocab 30522) with a tiny-random model (1124-row embedding) # feeds out-of-range token ids into gather and relies on backend # OOB behaviour, which is how we originally hit a :nan score. - @tag :distilbert_full - test "batched_run drives DistilBERT-QA through Nx.Serving" do + # `tag: :distilbert_full, lane_tags: false` keeps all three lanes gated + # behind `:distilbert_full` (not the lightweight `:native`), so they run + # under `--only distilbert_full` and never bloat `--only native`. The + # forward is driven through `Nx.Serving`'s `:defn_options`, which is + # where the compiler lanes plug in. + mode_test "batched_run drives DistilBERT-QA through Nx.Serving", + tag: :distilbert_full, + lane_tags: false do {:ok, model_info} = Bumblebee.load_model({:hf, "distilbert-base-uncased-distilled-squad"}) {:ok, tokenizer} = Bumblebee.load_tokenizer({:hf, "distilbert-base-uncased-distilled-squad"}) - serving = Bumblebee.Text.question_answering(model_info, tokenizer) + serving = + Bumblebee.Text.question_answering(model_info, tokenizer, defn_options: predict_opts) start_supervised!({Nx.Serving, serving: serving, name: __MODULE__.Serving}) diff --git a/test/emily/conformance/modernbert_test.exs b/test/emily/conformance/modernbert_test.exs index 99f9b40..b6ef5f1 100644 --- a/test/emily/conformance/modernbert_test.exs +++ b/test/emily/conformance/modernbert_test.exs @@ -33,7 +33,7 @@ defmodule Emily.Conformance.ModernBertTest do @moduletag :conformance @moduletag capture_log: true - test "ModernBert :base forward on Emily.Backend" do + mode_test "ModernBert :base forward on Emily.Backend" do spec = Bumblebee.configure(ModernBert, architecture: :base, @@ -47,7 +47,10 @@ defmodule Emily.Conformance.ModernBertTest do ) model = ModernBert.model(spec) - {init_fn, predict_fn} = Axon.build(model) + # Init on the evaluator (params are random-init, mode-irrelevant); + # gate only the forward pass under `predict_opts`. + {init_fn, _} = Axon.build(model) + {_, predict_fn} = Axon.build(model, predict_opts) input_template = %{ "input_ids" => Nx.template({1, 8}, :s64), diff --git a/test/emily/conformance/nomic_embeddings_test.exs b/test/emily/conformance/nomic_embeddings_test.exs index 7e74c80..4003f7e 100644 --- a/test/emily/conformance/nomic_embeddings_test.exs +++ b/test/emily/conformance/nomic_embeddings_test.exs @@ -27,7 +27,7 @@ defmodule Emily.Conformance.NomicEmbeddingsTest do @moduletag :conformance @moduletag capture_log: true - test "NomicBert :base forward runs end-to-end on Emily.Backend" do + mode_test "NomicBert :base forward runs end-to-end on Emily.Backend" do spec = Bumblebee.configure(NomicBert, architecture: :base, @@ -41,7 +41,10 @@ defmodule Emily.Conformance.NomicEmbeddingsTest do ) model = NomicBert.model(spec) - {init_fn, predict_fn} = Axon.build(model) + # Init on the evaluator (params are random-init, mode-irrelevant); + # gate only the forward pass under `predict_opts`. + {init_fn, _} = Axon.build(model) + {_, predict_fn} = Axon.build(model, predict_opts) input_template = %{ "input_ids" => Nx.template({1, 8}, :s64), diff --git a/test/emily/conformance/smollm3_test.exs b/test/emily/conformance/smollm3_test.exs index fdc7002..fca1179 100644 --- a/test/emily/conformance/smollm3_test.exs +++ b/test/emily/conformance/smollm3_test.exs @@ -30,7 +30,7 @@ defmodule Emily.Conformance.SmolLm3Test do @moduletag :conformance @moduletag capture_log: true - test "SmolLm3 :for_causal_language_modeling forward on Emily.Backend" do + mode_test "SmolLm3 :for_causal_language_modeling forward on Emily.Backend" do spec = Bumblebee.configure(SmolLm3, architecture: :for_causal_language_modeling, @@ -44,7 +44,10 @@ defmodule Emily.Conformance.SmolLm3Test do ) model = SmolLm3.model(spec) - {init_fn, predict_fn} = Axon.build(model) + # Init on the evaluator (params are random-init, mode-irrelevant); + # gate only the forward pass under `predict_opts`. + {init_fn, _} = Axon.build(model) + {_, predict_fn} = Axon.build(model, predict_opts) input_template = %{"input_ids" => Nx.template({1, 8}, :s64)} params = init_fn.(input_template, Axon.ModelState.empty()) diff --git a/test/emily/conformance/vit_full_test.exs b/test/emily/conformance/vit_full_test.exs index accb739..eabbb1e 100644 --- a/test/emily/conformance/vit_full_test.exs +++ b/test/emily/conformance/vit_full_test.exs @@ -30,7 +30,8 @@ defmodule Emily.Conformance.VitFullTest do @moduletag capture_log: true @moduletag timeout: 600_000 - test "google/vit-base-patch16-224 forward pass matches pinned logits slice" do + mode_test "google/vit-base-patch16-224 forward pass matches pinned logits slice", + lane_tags: false do {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "google/vit-base-patch16-224"}) @@ -42,7 +43,7 @@ defmodule Emily.Conformance.VitFullTest do "pixel_values" => Nx.broadcast(Nx.tensor(0.5, type: :f32), {1, 224, 224, 3}) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 1000} diff --git a/test/emily/conformance/vit_test.exs b/test/emily/conformance/vit_test.exs index baa39f7..5e38938 100644 --- a/test/emily/conformance/vit_test.exs +++ b/test/emily/conformance/vit_test.exs @@ -31,7 +31,7 @@ defmodule Emily.Conformance.VitTest do @moduletag capture_log: true @moduletag timeout: 120_000 - test ":base" do + mode_test ":base" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-ViTModel"}) @@ -41,7 +41,7 @@ defmodule Emily.Conformance.VitTest do "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.hidden_state) == {1, 226, 32} assert Nx.shape(outputs.pooled_state) == {1, 32} @@ -59,7 +59,7 @@ defmodule Emily.Conformance.VitTest do ) end - test ":for_image_classification" do + mode_test ":for_image_classification" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-ViTForImageClassification"} @@ -71,7 +71,7 @@ defmodule Emily.Conformance.VitTest do "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 2} @@ -81,7 +81,7 @@ defmodule Emily.Conformance.VitTest do ) end - test ":for_masked_image_modeling" do + mode_test ":for_masked_image_modeling" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-ViTForMaskedImageModeling"} @@ -93,7 +93,7 @@ defmodule Emily.Conformance.VitTest do "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.pixel_values) == {1, 30, 30, 3} diff --git a/test/emily/conformance/whisper_full_test.exs b/test/emily/conformance/whisper_full_test.exs index 78e48c5..30c3ce3 100644 --- a/test/emily/conformance/whisper_full_test.exs +++ b/test/emily/conformance/whisper_full_test.exs @@ -31,7 +31,7 @@ defmodule Emily.Conformance.WhisperFullTest do @moduletag capture_log: true @moduletag timeout: 600_000 - test "openai/whisper-tiny forward pass matches pinned logits slice" do + mode_test "openai/whisper-tiny forward pass matches pinned logits slice", lane_tags: false do {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "openai/whisper-tiny"}) @@ -58,7 +58,7 @@ defmodule Emily.Conformance.WhisperFullTest do "decoder_attention_mask" => decoder_attention_mask } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 6, 51_865} diff --git a/test/emily/conformance/whisper_test.exs b/test/emily/conformance/whisper_test.exs index 57f7d15..4bfb3ff 100644 --- a/test/emily/conformance/whisper_test.exs +++ b/test/emily/conformance/whisper_test.exs @@ -31,7 +31,7 @@ defmodule Emily.Conformance.WhisperTest do @moduletag capture_log: true @moduletag timeout: 300_000 - test ":base" do + mode_test ":base" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-WhisperModel"}) @@ -43,7 +43,7 @@ defmodule Emily.Conformance.WhisperTest do "decoder_attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.hidden_state) == {1, 8, 16} @@ -55,7 +55,7 @@ defmodule Emily.Conformance.WhisperTest do ) end - test ":for_conditional_generation" do + mode_test ":for_conditional_generation" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model( {:hf, "hf-internal-testing/tiny-random-WhisperForConditionalGeneration"} @@ -69,7 +69,7 @@ defmodule Emily.Conformance.WhisperTest do "decoder_attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 0, 0]]) } - outputs = Axon.predict(model, params, inputs) + outputs = Axon.predict(model, params, inputs, predict_opts) assert Nx.shape(outputs.logits) == {1, 8, 50_257} diff --git a/test/emily/opcode_parity_test.exs b/test/emily/opcode_parity_test.exs new file mode 100644 index 0000000..d7231a5 --- /dev/null +++ b/test/emily/opcode_parity_test.exs @@ -0,0 +1,67 @@ +defmodule Emily.OpcodeParityTest do + @moduledoc """ + Guards the one hand-maintained lockstep in the native compiler: the + opcode wire values live in **two** places that must agree — + `Emily.IR`'s `@opcodes` map (Elixir) and the `Opcode` enum + + `kOpcodeCount` in `c_src/emily/opcodes.hpp` (C++). Nothing in the build + enforces it: a mismatch compiles fine and only misbehaves at runtime + (an instruction dispatches to the wrong MLX op, or `valid_opcode` + rejects a real one). + + The check is value-based rather than name-based — both sides must be a + contiguous `0..N-1` with N == `kOpcodeCount`. That catches every + realistic drift (an opcode added to one side only, a forgotten + `kOpcodeCount` bump, a duplicate or gapped number) without depending on + the Elixir snake_case names matching the C++ PascalCase ones (they + don't always: `negate`/`Negative`, `fast_rms_norm`/`FastRMSNorm`). A + name/value *permutation* that keeps both contiguous would slip past + here, but that produces wrong results and is caught by + `Emily.CompilerEquivalenceTest`. + """ + use ExUnit.Case, async: true + + alias Emily.IR + + @header Path.expand("../../c_src/emily/opcodes.hpp", __DIR__) + + # Parse `kOpcodeCount` and the `Opcode` enum's explicit `Name = N,` + # values out of the C++ header. + defp parse_header do + src = File.read!(@header) + + [_, count] = Regex.run(~r/kOpcodeCount\s*=\s*(\d+)/, src) + count = String.to_integer(count) + + [_, body] = Regex.run(~r/enum class Opcode\s*:\s*int64_t\s*\{(.*?)\};/s, src) + + values = + ~r/^\s*[A-Za-z]\w*\s*=\s*(\d+)\s*,/m + |> Regex.scan(body) + |> Enum.map(fn [_, v] -> String.to_integer(v) end) + + %{count: count, enum_values: values} + end + + test "C++ Opcode enum is contiguous 0..N-1 and matches kOpcodeCount" do + %{count: count, enum_values: values} = parse_header() + + assert length(values) == count, + "opcodes.hpp has #{length(values)} enum entries but kOpcodeCount is #{count} — " <> + "bump kOpcodeCount when adding an opcode" + + assert Enum.sort(values) == Enum.to_list(0..(count - 1)), + "Opcode enum values are not a unique, gap-free 0..#{count - 1}" + end + + test "Emily.IR @opcodes stays in lockstep with the C++ enum" do + %{count: count} = parse_header() + opcodes = IR.opcodes() + + assert map_size(opcodes) == count, + "Emily.IR has #{map_size(opcodes)} opcodes but c_src/emily/opcodes.hpp kOpcodeCount " <> + "is #{count} — the two must be updated together" + + assert opcodes |> Map.values() |> Enum.sort() == Enum.to_list(0..(count - 1)), + "Emily.IR @opcodes values are not a unique, gap-free 0..#{count - 1}" + end +end diff --git a/test/emily/training/cnn_native_curve_test.exs b/test/emily/training/cnn_native_curve_test.exs new file mode 100644 index 0000000..f2d043a --- /dev/null +++ b/test/emily/training/cnn_native_curve_test.exs @@ -0,0 +1,134 @@ +defmodule Emily.Training.CnnNativeCurveTest do + @moduledoc """ + Native single-NIF CNN training convergence (issue #174). + + The training analogue of the conformance native lanes + (`Emily.Conformance.CompilerNativeTest`): a full conv + maxpool + training step — forward, backward, grad, and SGD update — is driven + through `compiler: Emily.Compiler, native: true, native_fallback: + :raise` for 30 steps and the per-step loss trajectory is checked + against two references. + + Why this exists. `cnn_curve_test.exs` already curve-matches the + handwritten CNN, but only in **eval** mode (`Emily.Compiler` walking + the Expr op-by-op via the Evaluator). Every other `training/*` test + is eval-only too, so CNN training was verified-lowering (the + `compiler_equivalence_test.exs` op gates) but never **convergence**- + tested under the single-NIF replay. This closes that gap. + + Three lanes, same deterministic init and data: + + * **native** — `native: true, native_fallback: :raise`. The + `:raise` makes this a no-fallback gate: if any op in the + forward+backward+grad+SGD step fails to lower (the maxpool + backward lands on `window_scatter_max` every step; the conv + backward flips the kernel with `reverse`), the run raises here + instead of silently degrading to the evaluator. + * **eval** — `Emily.Compiler` op-by-op. Same MLX kernels in the + same order as the native replay, so the two track **bit- + identically** through training. A 1e-6 bar asserts the single- + NIF lowering reproduces op-by-op exactly across 30 SGD updates. + * **binary** — `Nx.Defn.Evaluator` on `Nx.BinaryBackend`, the + non-MLX convergence oracle. Looser bar (1e-2 rtol, as in + `cnn_curve_test.exs`) absorbs f32 reduction-order drift between + MLX's parallel reductions and BinaryBackend's sequential ones. + + No Axon — the handwritten path keeps the failure surface tiny (see + `cnn_curve_test.exs`). The Axon CNN canary stays in + `mnist_cnn_full_test.exs` (`:training_full`). + """ + + use ExUnit.Case, async: true + + alias Emily.TrainingHelper, as: TH + import TH, only: [close?: 4, flunk_trajectory: 5] + + @native [compiler: Emily.Compiler, native: true, native_fallback: :raise] + @eval [compiler: Emily.Compiler] + + @input_shape {1, 10, 10} + @batch 4 + @classes 3 + @steps 30 + @lr_val 0.05 + + test "per-step CNN loss trajectory matches under native single-NIF compile" do + # Native single-NIF lane — the system under test. `native_fallback: + # :raise` proves full native coverage of the training step. + params_native = TH.init_cnn(@input_shape, @classes, 0, Emily.Backend) + {x_native, y_native} = TH.cnn_batch({@batch, 10, 10}, @classes, Emily.Backend) + lr_native = Nx.tensor(@lr_val, type: {:f, 32}, backend: Emily.Backend) + + losses_native = + TH.run_steps( + &TH.cnn_step_with_loss/4, + params_native, + [x_native, y_native, lr_native], + @steps, + @native + ) + + # Op-by-op Emily eval lane — same MLX kernels, isolates single-NIF + # lowering bugs from backend numerics. + params_eval = TH.init_cnn(@input_shape, @classes, 0, Emily.Backend) + {x_eval, y_eval} = TH.cnn_batch({@batch, 10, 10}, @classes, Emily.Backend) + lr_eval = Nx.tensor(@lr_val, type: {:f, 32}, backend: Emily.Backend) + + losses_eval = + TH.run_steps( + &TH.cnn_step_with_loss/4, + params_eval, + [x_eval, y_eval, lr_eval], + @steps, + @eval + ) + + # BinaryBackend oracle — the non-MLX convergence reference. + params_bin = TH.init_cnn(@input_shape, @classes, 0, Nx.BinaryBackend) + {x_bin, y_bin} = TH.cnn_batch({@batch, 10, 10}, @classes, Nx.BinaryBackend) + lr_bin = Nx.tensor(@lr_val, type: {:f, 32}, backend: Nx.BinaryBackend) + + losses_bin = + TH.run_steps( + &TH.cnn_step_with_loss/4, + params_bin, + [x_bin, y_bin, lr_bin], + @steps, + Nx.Defn.Evaluator + ) + + assert length(losses_native) == @steps + assert length(losses_eval) == @steps + assert length(losses_bin) == @steps + + # 1. Single-NIF native == op-by-op eval. Both are MLX in the same + # order, so they track bit-identically; the tight bar makes a + # divergent native trajectory a hard failure. + for {{ln, le}, i} <- Enum.zip(losses_native, losses_eval) |> Enum.with_index() do + close?(ln, le, 1.0e-6, 1.0e-6) || + flunk_trajectory(i, ln, le, losses_native, losses_eval) + end + + # 2. Native trajectory matches the BinaryBackend oracle within the + # CNN tolerance — same bar as cnn_curve_test.exs. + for {{ln, lb}, i} <- Enum.zip(losses_native, losses_bin) |> Enum.with_index() do + close?(ln, lb, 1.0e-4, 1.0e-2) || + flunk_trajectory(i, ln, lb, losses_native, losses_bin) + end + + # 3. Convergence — the native loss actually decreased over the run. + assert List.first(losses_native) > List.last(losses_native), + "native loss did not decrease: first=#{List.first(losses_native)} " <> + "last=#{List.last(losses_native)}" + + # 4. Final loss agrees with the oracle (convergence correctness: + # catches a run where per-step drift averaged out but the + # optimizer ended up somewhere wrong). + ln_final = List.last(losses_native) + lb_final = List.last(losses_bin) + + assert close?(ln_final, lb_final, 1.0e-4, 1.0e-2), + "final loss divergence: native=#{ln_final} bin=#{lb_final} " <> + "reldiff=#{abs(ln_final - lb_final) / abs(lb_final)}" + end +end diff --git a/test/emily/training/mlp_native_curve_test.exs b/test/emily/training/mlp_native_curve_test.exs new file mode 100644 index 0000000..3aca951 --- /dev/null +++ b/test/emily/training/mlp_native_curve_test.exs @@ -0,0 +1,113 @@ +defmodule Emily.Training.MlpNativeCurveTest do + @moduledoc """ + Native single-NIF MLP training convergence (issue #174). + + The dense + SGD companion to `cnn_native_curve_test.exs`: a 2-layer + ReLU MLP training step — forward, backward, grad, SGD update — is + driven through `compiler: Emily.Compiler, native: true, + native_fallback: :raise` for 50 steps and the per-step loss + trajectory is checked against two references. + + This closes the matmul-dominated half of the training-coverage gap + the issue calls out: `mlp_curve_test.exs` already curve-matches this + MLP, but only in eval mode. Here the same step replays through the + single NIF, with `:raise` proving the dense forward/backward and SGD + update lower with **zero** fallback. + + Three lanes, same deterministic init and data: + + * **native** — single-NIF replay, no-fallback gate. + * **eval** — `Emily.Compiler` op-by-op; bit-identical to native + (same MLX kernels, same order), asserted at a 1e-6 bar. + * **binary** — `Nx.Defn.Evaluator` on `Nx.BinaryBackend`, the + non-MLX oracle. The MLP is matmul-dominated, so the bar matches + `mlp_curve_test.exs` (1e-3 per-step rtol, 1e-4 final). + """ + + use ExUnit.Case, async: true + + alias Emily.TrainingHelper, as: TH + import TH, only: [close?: 4, flunk_trajectory: 5] + + @native [compiler: Emily.Compiler, native: true, native_fallback: :raise] + @eval [compiler: Emily.Compiler] + + @dims {4, 8, 3} + @batch_shape {16, 4, 3} + @steps 50 + @lr_val 0.5 + + test "per-step MLP loss trajectory matches under native single-NIF compile" do + # Native single-NIF lane — the system under test. + params_native = TH.init_mlp(@dims, 0, Emily.Backend) + {x_native, y_native} = TH.mlp_batch(@batch_shape, Emily.Backend) + lr_native = Nx.tensor(@lr_val, type: {:f, 32}, backend: Emily.Backend) + + losses_native = + TH.run_steps( + &TH.mlp_step_with_loss/4, + params_native, + [x_native, y_native, lr_native], + @steps, + @native + ) + + # Op-by-op Emily eval lane — same MLX kernels. + params_eval = TH.init_mlp(@dims, 0, Emily.Backend) + {x_eval, y_eval} = TH.mlp_batch(@batch_shape, Emily.Backend) + lr_eval = Nx.tensor(@lr_val, type: {:f, 32}, backend: Emily.Backend) + + losses_eval = + TH.run_steps( + &TH.mlp_step_with_loss/4, + params_eval, + [x_eval, y_eval, lr_eval], + @steps, + @eval + ) + + # BinaryBackend oracle. + params_bin = TH.init_mlp(@dims, 0, Nx.BinaryBackend) + {x_bin, y_bin} = TH.mlp_batch(@batch_shape, Nx.BinaryBackend) + lr_bin = Nx.tensor(@lr_val, type: {:f, 32}, backend: Nx.BinaryBackend) + + losses_bin = + TH.run_steps( + &TH.mlp_step_with_loss/4, + params_bin, + [x_bin, y_bin, lr_bin], + @steps, + Nx.Defn.Evaluator + ) + + assert length(losses_native) == @steps + assert length(losses_eval) == @steps + assert length(losses_bin) == @steps + + # 1. Single-NIF native == op-by-op eval (bit-identical MLX path). + for {{ln, le}, i} <- Enum.zip(losses_native, losses_eval) |> Enum.with_index() do + close?(ln, le, 1.0e-6, 1.0e-6) || + flunk_trajectory(i, ln, le, losses_native, losses_eval) + end + + # 2. Native trajectory matches the BinaryBackend oracle — same bar + # as mlp_curve_test.exs. + for {{ln, lb}, i} <- Enum.zip(losses_native, losses_bin) |> Enum.with_index() do + close?(ln, lb, 1.0e-4, 1.0e-3) || + flunk_trajectory(i, ln, lb, losses_native, losses_bin) + end + + # 3. Convergence — the native loss actually decreased over the run. + assert List.first(losses_native) > List.last(losses_native), + "native loss did not decrease: first=#{List.first(losses_native)} " <> + "last=#{List.last(losses_native)}" + + # 4. Final loss agrees with the oracle (convergence correctness). + ln_final = List.last(losses_native) + lb_final = List.last(losses_bin) + + assert close?(ln_final, lb_final, 1.0e-5, 1.0e-4), + "final loss divergence: native=#{ln_final} bin=#{lb_final} " <> + "reldiff=#{abs(ln_final - lb_final) / abs(lb_final)}" + end +end diff --git a/test/emily/training/mnist_cnn_native_full_test.exs b/test/emily/training/mnist_cnn_native_full_test.exs new file mode 100644 index 0000000..ed2e07c --- /dev/null +++ b/test/emily/training/mnist_cnn_native_full_test.exs @@ -0,0 +1,84 @@ +defmodule Emily.Training.MnistCnnNativeFullTest do + @moduledoc """ + MNIST CNN convergence under the **native single-NIF compiler** + (issue #174, `:training_full`). + + The native-lane analogue of `mnist_cnn_full_test.exs`: the same + LeNet-style Axon CNN trains on real MNIST, but the whole training + step compiles through `compiler: Emily.Compiler, native: true, + native_fallback: :raise` instead of the op-by-op evaluator. + + Why `native_fallback: :raise` makes this self-proving. Once + `native: true` reaches `Emily.Compiler`, the run is binary: the Expr + lowers to one program and replays in a single NIF, or an op it can't + lower raises (`Emily.Compiler.build_native/4` reraises the lowering + `ArgumentError` under `:raise` — it can never silently degrade to the + evaluator, that path only exists under `:eval`). So a training run + that *completes* proves the entire step — forward (conv, ReLU, + maxpool), categorical-cross-entropy loss, the backward + (`window_scatter_max` for the maxpool grad, `reverse` for the conv + kernel flip), and the Adam update — lowered fully native with zero + fallback. `Axon.Loop.run` forwards the per-call jit options (it pops + only `:jit_compile?`/`:force_garbage_collection?`), so the options + genuinely reach the compiler. + + Native replay is bit-identical to the evaluator (same MLX kernels in + the same order), so the accuracy bar matches the eval canary exactly + (>97%). The point isn't a different number — it's that training + reaches it through the single-NIF path. + + Opt-in — `mix test --only training_full` (downloads MNIST, multi- + minute training). + """ + + use ExUnit.Case, async: true + + alias Emily.MnistHelper + + @moduletag :training_full + @moduletag capture_log: true + @moduletag timeout: 600_000 + + setup do + Nx.default_backend(Emily.Backend) + :ok + end + + @batch_size 64 + @epochs 5 + @target_accuracy 0.97 + + # Strict no-fallback native lane. `:raise` makes a completed run a + # proof of full native lowering (see the moduledoc). + @native [compiler: Emily.Compiler, native: true, native_fallback: :raise] + + test "Axon CNN reaches >#{trunc(@target_accuracy * 100)}% accuracy via the native single-NIF compiler" do + {train_batches, test_images, test_labels} = MnistHelper.load_mnist(@batch_size, :cnn) + + # Channels-last (Axon default) — MnistHelper produces {N, 28, 28, 1}. + model = + Axon.input("input", shape: {nil, 28, 28, 1}) + |> Axon.conv(8, kernel_size: {3, 3}, activation: :relu) + |> Axon.max_pool(kernel_size: {2, 2}, strides: [2, 2]) + |> Axon.conv(16, kernel_size: {3, 3}, activation: :relu) + |> Axon.max_pool(kernel_size: {2, 2}, strides: [2, 2]) + |> Axon.flatten() + |> Axon.dense(64, activation: :relu) + |> Axon.dense(10, activation: :softmax) + + # The whole training loop (init + step) compiles native: `Axon.Loop.run` + # forwards `native:`/`native_fallback:` to the defn jit. Under `:raise`, + # reaching the end proves every op lowered — no silent eval fallback. + trained_state = + model + |> Axon.Loop.trainer(:categorical_cross_entropy, :adam) + |> Axon.Loop.run(train_batches, %{}, [epochs: @epochs] ++ @native) + + # Evaluate through the native path too, so the accuracy that gates the + # test is itself produced by the single-NIF forward. + accuracy = MnistHelper.evaluate(model, trained_state, test_images, test_labels, @native) + + assert accuracy >= @target_accuracy, + "native MNIST CNN accuracy #{Float.round(accuracy, 4)} below target #{@target_accuracy}" + end +end diff --git a/test/emily/training/mnist_native_full_test.exs b/test/emily/training/mnist_native_full_test.exs new file mode 100644 index 0000000..dc50259 --- /dev/null +++ b/test/emily/training/mnist_native_full_test.exs @@ -0,0 +1,62 @@ +defmodule Emily.Training.MnistNativeFullTest do + @moduledoc """ + MNIST MLP convergence under the **native single-NIF compiler** + (issue #174, `:training_full`). + + The native-lane analogue of `mnist_full_test.exs`: the same dense + MLP trains on real MNIST, but the whole training step compiles + through `compiler: Emily.Compiler, native: true, native_fallback: + :raise` instead of the op-by-op evaluator. Closes the matmul- + dominated (dense + Adam) half of the training-coverage gap; the + conv/pooling half is `mnist_cnn_native_full_test.exs`. + + `native_fallback: :raise` makes the run self-proving: once + `native: true` reaches `Emily.Compiler` the outcome is binary — + the step lowers to one program and replays in a single NIF, or an + un-lowerable op raises (it can never silently degrade to the + evaluator under `:raise`). So reaching the accuracy assertion proves + the dense forward, cross-entropy loss, backward, and Adam update all + lowered fully native. See `mnist_cnn_native_full_test.exs` for the + full rationale. + + Opt-in — `mix test --only training_full`. + """ + + use ExUnit.Case, async: true + + alias Emily.MnistHelper + + @moduletag :training_full + @moduletag capture_log: true + @moduletag timeout: 600_000 + + setup do + Nx.default_backend(Emily.Backend) + :ok + end + + @batch_size 128 + @epochs 5 + @target_accuracy 0.96 + + @native [compiler: Emily.Compiler, native: true, native_fallback: :raise] + + test "Axon MLP reaches >#{trunc(@target_accuracy * 100)}% accuracy via the native single-NIF compiler" do + {train_batches, test_images, test_labels} = MnistHelper.load_mnist(@batch_size) + + model = + Axon.input("input", shape: {nil, 784}) + |> Axon.dense(128, activation: :relu) + |> Axon.dense(10, activation: :softmax) + + trained_state = + model + |> Axon.Loop.trainer(:categorical_cross_entropy, :adam) + |> Axon.Loop.run(train_batches, %{}, [epochs: @epochs] ++ @native) + + accuracy = MnistHelper.evaluate(model, trained_state, test_images, test_labels, @native) + + assert accuracy >= @target_accuracy, + "native MNIST MLP accuracy #{Float.round(accuracy, 4)} below target #{@target_accuracy}" + end +end diff --git a/test/support/conformance_helper.ex b/test/support/conformance_helper.ex index 0ec9a52..51ddbad 100644 --- a/test/support/conformance_helper.ex +++ b/test/support/conformance_helper.ex @@ -29,7 +29,8 @@ defmodule Emily.ConformanceHelper do defmacro __using__(_opts) do quote do - import Emily.ConformanceHelper, only: [assert_all_close: 2, assert_all_close: 3] + import Emily.ConformanceHelper, + only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2, mode_test: 3] setup do Nx.default_backend(Emily.Backend) @@ -38,6 +39,119 @@ defmodule Emily.ConformanceHelper do end end + @doc """ + Define a conformance test in three lanes from a single body. + + Expands to three `ExUnit` tests that share `body` but bind a different + `predict_opts` keyword list: + + * the default lane binds `predict_opts` to `[]` — the evaluator path + Bumblebee/Axon use out of the box (the existing "eval'd" mode); + * the native lane binds `predict_opts` to + `[compiler: Emily.Compiler, native: true, native_fallback: :raise]` + and is additionally tagged `:native`; + * the fusion lane adds `native_compiled: true` (wrapping the replay in + `mx::compile`) and is additionally tagged `:native_compiled`. + + The module is already tagged `:conformance`, so the native and fusion + lanes carry that tag too: `mix test --only conformance` runs all three, + while `mix test --only native` / `mix test --only native_compiled` run + one lane each. Because every lane resolves the same HuggingFace repos, + whichever runs first reads from `~/.cache/bumblebee` for the rest — the + download is paid once. + + `mx::compile` reassociates f32, so the fusion lane's logits are not + bit-identical to the evaluator's; it shares the same reference and + tolerance as the other lanes (these tiny-random forwards drift well + within `assert_all_close`'s default), and `assert_finite!`-style smoke + tests are robust to the drift outright. + + The body must thread `predict_opts` into whatever drives the forward + pass so the two lanes assert against the *identical* reference and + cannot drift apart in maintenance: + + mode_test ":base" do + {:ok, %{model: model, params: params}} = Bumblebee.load_model(...) + outputs = Axon.predict(model, params, inputs, predict_opts) + assert_all_close(outputs.hidden_state, ...) + end + + For `Axon.build`-driven tests, build `init_fn` on the evaluator (params + are random-init, mode-irrelevant) and only `predict_fn` under + `predict_opts`, so the native lane gates the forward pass alone: + + {init_fn, _} = Axon.build(model) + {_, predict_fn} = Axon.build(model, predict_opts) + + `native_fallback: :raise` makes the native lane a no-fallback gate: an + op that does not lower fails the test rather than silently degrading to + the evaluator, so a red native lane is a concrete op-coverage gap. + + ## Options + + * `:lane_tags` (default `true`) — when `false`, the native and fusion + lanes are emitted *without* the cross-cutting `:native` / + `:native_compiled` tags. The heavyweight `*_full` suites pass + `lane_tags: false` so their compiler lanes stay gated behind the + suite's own `:*_full` moduletag; otherwise `--only native` would + start pulling full-size checkpoints. `--only vit_full` then runs all + three lanes of that suite. + + * `:tag` — an extra tag stamped on *every* lane. Used by the + `Nx.Serving` test, which lives in a `:conformance`-tagged module but + must stay gated behind `:distilbert_full` like its eval lane: + `tag: :distilbert_full, lane_tags: false`. + """ + defmacro mode_test(name, opts \\ [], do: body) do + tag_lanes? = Keyword.get(opts, :lane_tags, true) + extra_tag = Keyword.get(opts, :tag) + + lanes = [ + lane([extra_tag], name, "", [], body), + lane( + [extra_tag, tag_lanes? && :native], + name, + " [native]", + [compiler: Emily.Compiler, native: true, native_fallback: :raise], + body + ), + lane( + [extra_tag, tag_lanes? && :native_compiled], + name, + " [native_compiled]", + [compiler: Emily.Compiler, native: true, native_fallback: :raise, native_compiled: true], + body + ) + ] + + quote do + (unquote_splicing(lanes)) + end + end + + # Build one `mode_test` lane: a `test` that binds `predict_opts` for the + # body, preceded by one `@tag` per entry in `tags` (nil/false entries are + # dropped). The `*_full` suites pass `lane_tags: false` to drop the + # `:native` / `:native_compiled` tags and rely on their own `:*_full` + # moduletag (or an explicit `:tag`) instead. + defp lane(tags, name, suffix, predict_opts, body) do + tags = Enum.reject(tags, &(&1 in [nil, false])) + + name_ast = + if suffix == "", do: name, else: quote(do: unquote(name) <> unquote(suffix)) + + tag_attrs = for t <- tags, do: quote(do: @tag(unquote(t))) + + quote do + (unquote_splicing(tag_attrs)) + + test unquote(name_ast) do + var!(predict_opts) = unquote(predict_opts) + unquote(body) + end + end + end + @doc """ Assert that every element of `left` agrees with `right` within `atol + rtol * |right|`. diff --git a/test/support/mnist_helper.ex b/test/support/mnist_helper.ex index 730c6c4..c54a260 100644 --- a/test/support/mnist_helper.ex +++ b/test/support/mnist_helper.ex @@ -24,9 +24,13 @@ defmodule Emily.MnistHelper do {train_batches, test_images, test_labels} end - def evaluate(model, state, test_images, test_labels) do + # `predict_opts` are forwarded to `Axon.predict` (which forwards them to + # the defn jit). Defaults to the op-by-op eval lane; the native-lane + # tests pass `[compiler: Emily.Compiler, native: true, native_fallback: + # :raise]` so the gating accuracy is itself produced by the single NIF. + def evaluate(model, state, test_images, test_labels, predict_opts \\ [compiler: Emily.Compiler]) do logits = - Axon.predict(model, state, test_images, compiler: Emily.Compiler) + Axon.predict(model, state, test_images, predict_opts) predicted = Nx.argmax(logits, axis: -1) actual = Nx.argmax(test_labels, axis: -1) diff --git a/test/support/training_helper.ex b/test/support/training_helper.ex index e683762..1005454 100644 --- a/test/support/training_helper.ex +++ b/test/support/training_helper.ex @@ -262,12 +262,20 @@ defmodule Emily.TrainingHelper do `args` is a list of the remaining non-params tensors passed after params on each call — e.g. `[x, y, lr]` for MLP, `[x, y, lr, scale]` for the transformer block. + + The final argument is either a bare compiler module (e.g. + `Emily.Compiler`, `Nx.Defn.Evaluator`) or a full `Nx.Defn.jit` opts + keyword list. The latter is how the native single-NIF lane is driven: + + run_steps(fun, params, args, n, + compiler: Emily.Compiler, native: true, native_fallback: :raise) """ - def run_steps(step_fun, params, args, n, compiler) when is_list(args) do + def run_steps(step_fun, params, args, n, compiler_or_opts) when is_list(args) do + opts = step_opts(compiler_or_opts) + {_final, losses_rev} = Enum.reduce(1..n, {params, []}, fn _i, {params, losses} -> - {new_params, loss} = - Nx.Defn.jit_apply(step_fun, [params | args], compiler: compiler) + {new_params, loss} = Nx.Defn.jit_apply(step_fun, [params | args], opts) loss_f = loss |> Nx.backend_transfer(Nx.BinaryBackend) |> Nx.to_number() {new_params, [loss_f | losses]} @@ -276,6 +284,12 @@ defmodule Emily.TrainingHelper do Enum.reverse(losses_rev) end + # Accept either a bare compiler module (back-compat with the eval-lane + # callers) or a full jit opts list (the native lane passes `native:`/ + # `native_fallback:` through, which a bare `compiler:` can't carry). + defp step_opts(opts) when is_list(opts), do: opts + defp step_opts(compiler) when is_atom(compiler), do: [compiler: compiler] + # -------------------- Curve-matching assertions -------------------- @doc """ diff --git a/test/test_helper.exs b/test/test_helper.exs index 58b751d..2205023 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -41,10 +41,29 @@ # the MLX `mx::fast::*` kernels via `Emily.Fast`. Run explicitly: # # mix test --only fast_kernels_full +# +# `:native` and `:native_compiled` are the expression-compiler lanes of +# the tiny-random conformance suites: every `mode_test` (see +# `Emily.ConformanceHelper`) re-runs the forward pass under +# `compiler: Emily.Compiler, native: true, native_fallback: :raise` +# (`:native`) and again with `native_compiled: true` wrapping the replay +# in `mx::compile` (`:native_compiled`), so the same PyTorch reference +# slice validates the evaluator, the native-compiled, and the fused +# paths. Those tests carry `:conformance` too, so `--only conformance` +# runs all three lanes; select one lane alone with: +# +# mix test --only native +# mix test --only native_compiled +# +# Listed in the default exclude defensively — every such test is already +# `:conformance`-tagged, but this keeps a future `:native`-only or +# `:native_compiled`-only test out of the default suite. ExUnit.start( max_cases: System.schedulers_online(), exclude: [ :conformance, + :native, + :native_compiled, :vit_full, :whisper_full, :distilbert_full,