M17: native window ops (conv-pool training) - #34
Merged
Conversation
Lift window_sum, window_max, window_min, window_product, window_scatter_max, window_scatter_min off via_binary onto native MLX. MLX has no direct primitives — each op is composed as mx::pad → mx::as_strided (sliding-window view) → reduce, mirroring nn/layers/pooling.py generalised to N-D so Nx's per-axis window/stride/padding/dilation all lower cleanly. Scatter variants add argmax-with-last-occurrence-tie-break (mask * arange(K) trick, since mx::argmax is first-occurrence) → per-axis absolute-index computation → mx::scatter_add into a mx::full(init_value) buffer → slice back. Closes the training primitive set: grad(window_max) in Nx rewrites to window_scatter_max, so every MaxPool backward pass now stays on MLX instead of round-tripping through BinaryBackend. Tests: backend_window_test (19 unit — shape × stride × padding × dilation × f32/bf16/s32/u8), backend_window_scatter_test (11 — overlapping windows, tie-break, non-zero init, 1-D/3-D), grad-equivalence extensions (3 new zoo fns: window_sum, window_max_pool, window_avg_pool; bf16 equivalence auto-picks-up; EXLA oracle skips un-regenerated entries via ExlaGoldenData.has_golden?/1), cnn_curve_test (handwritten 2-conv + max-pool CNN, 30 steps), mnist_cnn_full_test (:training_full, LeNet ≥97%, observed 1.60 → 0.12 loss over 5 epochs).
Regenerate test/support/exla_golden_data.ex via bench/exla_golden_gen.exs with EXLA 0.11.0 CPU. Adds the three M17 window zoo functions (grad_window_sum, grad_window_max_pool, grad_window_avg_pool) to both the generator's @zoo and the oracle test's @Tolerance table, restoring full cross-backend coverage. Emily's last-occurrence argmax tie-break matches XLA's convention — all three pass the (1e-6, 1e-5) tolerance, tight enough to catch any numerics drift since the values are integer/rational (window counts / 1/4 averages / one-hot scatter). Removes the temporary ExlaGoldenData.has_golden?/1 skip scaffolding — no longer needed now that every zoo entry has a golden.
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.
Summary
window_sum,window_max,window_min,window_product,window_scatter_max,window_scatter_minoffvia_binaryonto native MLX — composed aspad → as_strided → reducefor the reductions,pad → as_strided → argmax-tie-break → scatter_addfor the scatters. Closes M17 in PLAN.md.window_maxin Nx rewrites towindow_scatter_max, so every MaxPool2d backward now stays on MLX instead of round-tripping through BinaryBackend — the stated M17 exit criterion ("small-CNN MNIST training converges in:training_full").c_src/ops/pooling.cpp; the Backend resolves:valid/:samepadding and dtype-specific identities (0/1/±∞, plus int-dtype min/max) in Elixir before dispatch.Test plan
mix precommit— 401 tests + 76 properties greenmix test test/emily/backend_window_test.exs— 19 forward-parity cases (shape × stride × padding × dilation × {f32, bf16, s32, u8})mix test test/emily/backend_window_scatter_test.exs— 11 cases incl. overlapping windows, last-occurrence tie-break, non-zero init, 1-D/3-Dmix test test/emily/grad/— grad-equivalence extensions: 3 new zoo fns (grad_window_sum,grad_window_max_pool,grad_window_avg_pool);bf16_grad_equivalence_testauto-picks-up; EXLA oracle covers all 11 zoo entriesmix test test/emily/training/cnn_curve_test.exs— handwritten 2-conv + max-pool CNN, 30-step loss trajectory match vs BinaryBackend within rtol 1e-2mix test --only training_full test/emily/training/mnist_cnn_full_test.exs— LeNet-style Axon CNN, 1.60 → 0.12 loss over 5 epochs, ≥ 97% test accuracy