Lower take_along_axis natively + opcode lockstep guard - #177
Merged
Conversation
Add a `mode_test/2` macro to `Emily.ConformanceHelper` that emits each conformance test in three lanes from a single body, varying only the bound `predict_opts`: * eval — `[]`, the evaluator path (unchanged behaviour) * native — `native: true, native_fallback: :raise` * fusion — additionally `native_compiled: true` (mx::compile) The native and fusion lanes are tagged `:native` and `:native_compiled` on top of the module's `:conformance`, so `--only conformance` runs all three while `--only native` / `--only native_compiled` select one lane. Both compiler lanes resolve the same HuggingFace repos as the evaluator lane, so the checkpoint download is amortised across them. Convert the six tiny-random suites (DistilBERT, ViT, Whisper, ModernBert, NomicBert, SmolLm3). `Axon.predict`-based tests thread `predict_opts` into the call; `Axon.build`-based smoke tests build `init_fn` on the evaluator and only `predict_fn` under `predict_opts`, so the native gate covers the forward pass rather than random param init. The serving and fast-kernels tests stay eval-only. `native_fallback: :raise` makes the native lanes a no-fallback op-coverage gate: every architecture lowers fully, and the fusion lane holds at the existing 1e-4 reference tolerance (mx::compile only reassociates f32, so its drift is far below the approximate mx::fast::* kernels).
Extend `mode_test/2` to `mode_test/3` with a `lane_tags:` option and run the `*_full` forward-pass conformance suites (ViT-Base, Whisper-tiny) through the native and fusion compilers alongside the evaluator. The full suites pass `lane_tags: false`, so their native and fusion lanes are emitted without the cross-cutting `:native` / `:native_compiled` tags and stay gated behind the suite's own `:vit_full` / `:whisper_full` moduletag. `--only vit_full` now runs all three lanes on the full checkpoint, while `--only native` stays tiny-random only and never pulls a full-size download. Both full forwards lower fully under `native_fallback: :raise` and hold at the pinned 1e-4 reference tolerance in the fusion lane too — `mx::compile` only reassociates f32, so its drift stays well below the approximate `mx::fast::*` kernels (which need 1e-3 ViT / 1e-2 Whisper in the `:fast_kernels_full` variants). The serving and fast-kernels tests remain eval-only.
`Nx.take_along_axis` (the `Nx.Block.TakeAlongAxis` block) now lowers under `compiler: Emily.Compiler, native: true` instead of raising. The MLX op and eager NIF already existed (`Emily.Native.take_along_axis` → `mlx::core::take_along_axis`, used by `Emily.Backend.native_take_along_axis/4`); only the IR path was missing. - ir.ex: add opcode `take_along_axis` (81) and a `lower_block/5` clause mirroring the eager twin — cast indices to s32, emit, coerce to output type — alongside the existing `Nx.Block.Take` lowering. - opcodes.hpp: add `TakeAlongAxis = 81`, bump `kOpcodeCount`, and a dispatch case calling `mx::take_along_axis(in[0], in[1], axis, s)`. - compiler_equivalence_test.exs: native-vs-evaluator bit-identical cases (last axis, axis 0, a 3-D transformer-shaped gather). This was the last op forcing a fallback in `Bumblebee.Text.question_answering`'s answer-span gather, so the DistilBERT-QA `Nx.Serving` forward now compiles fully native and fused. Wire it into the conformance suite: extend `mode_test` with a `:tag` option (stamp every lane) so the serving test runs its eval / native / native_compiled lanes gated behind `:distilbert_full`.
The opcode wire values are hand-maintained in two places — `Emily.IR`'s `@opcodes` map and the `Opcode` enum + `kOpcodeCount` in `c_src/emily/opcodes.hpp` — with nothing enforcing they agree. A mismatch compiles fine and only misbehaves at runtime. Add `Emily.OpcodeParityTest`, which parses the header and asserts both sides are a unique, gap-free `0..N-1` with N == `kOpcodeCount`. The check is value-based, not name-based, so it doesn't depend on the snake_case names matching the PascalCase enum (`negate`/`Negative`, `fast_rms_norm`/`FastRMSNorm`); a name/value permutation that keeps both contiguous is still caught by the equivalence suite. Expose `Emily.IR.opcodes/0` for the test to read the map.
ausimian
changed the base branch from
feat/expr-compiler-conformance-native-lanes
to
feat/expr-compiler
June 6, 2026 05:37
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
feat: lower take_along_axis in the native Expr compiler—Nx.take_along_axis(theNx.Block.TakeAlongAxisblock) now lowersunder
compiler: Emily.Compiler, native: trueinstead of raising.Adds opcode
take_along_axis(81) + alower_blockclause mirroringthe eager twin, the
opcodes.hppdispatch (mx::take_along_axis), andbit-identical native-vs-evaluator equivalence cases. This was the last
op forcing a fallback in
Bumblebee.Text.question_answering's answer-span gather, so the DistilBERT-QA serving forward now compiles fully
native (wired into the conformance suite via a
:tagoption).test: guard opcode lockstep between Emily.IR and the C++ enum—the opcode wire values are hand-maintained in
Emily.IR's@opcodesand the
Opcodeenum +kOpcodeCountinopcodes.hpp, with nothingenforcing agreement. Adds
Emily.OpcodeParityTest(value-based: bothsides a unique gap-free
0..N-1,N == kOpcodeCount) + exposesEmily.IR.opcodes/0.Pairs the new-opcode feature with the guard that protects opcode
additions. Builds on the conformance-helper changes in the lane PR.
PR 2/5. Stacked on
…-conformance-native-lanes; lands infeat/expr-compiler.