From c7d6f33d0d0da2a539be4a65e2d254d4dee661c8 Mon Sep 17 00:00:00 2001
From: ausimian <nick@ausimian.net>
Date: Sat, 6 Jun 2026 13:13:57 +1000
Subject: [PATCH 1/4] test: run conformance suites through the native compiler
 lanes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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).
---
 test/emily/conformance/distilbert_test.exs    | 26 +++---
 test/emily/conformance/modernbert_test.exs    |  7 +-
 .../conformance/nomic_embeddings_test.exs     |  7 +-
 test/emily/conformance/smollm3_test.exs       |  7 +-
 test/emily/conformance/vit_test.exs           | 12 +--
 test/emily/conformance/whisper_test.exs       |  8 +-
 test/support/conformance_helper.ex            | 83 ++++++++++++++++++-
 test/test_helper.exs                          | 19 +++++
 8 files changed, 139 insertions(+), 30 deletions(-)

diff --git a/test/emily/conformance/distilbert_test.exs b/test/emily/conformance/distilbert_test.exs
index 4929472..b590941 100644
--- a/test/emily/conformance/distilbert_test.exs
+++ b/test/emily/conformance/distilbert_test.exs
@@ -33,7 +33,7 @@ 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]
 
   alias Emily.Bumblebee.FastKernels
 
@@ -52,7 +52,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 +63,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 +75,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 +86,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 +98,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 +111,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 +131,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 +141,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 +154,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 +170,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 +183,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}
 
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_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_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/support/conformance_helper.ex b/test/support/conformance_helper.ex
index 0ec9a52..0e1e8d2 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]
 
       setup do
         Nx.default_backend(Emily.Backend)
@@ -38,6 +39,86 @@ 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.
+  """
+  defmacro mode_test(name, do: body) do
+    quote do
+      test unquote(name) do
+        var!(predict_opts) = []
+        unquote(body)
+      end
+
+      @tag :native
+      test unquote(name) <> " [native]" do
+        var!(predict_opts) = [
+          compiler: Emily.Compiler,
+          native: true,
+          native_fallback: :raise
+        ]
+
+        unquote(body)
+      end
+
+      @tag :native_compiled
+      test unquote(name) <> " [native_compiled]" do
+        var!(predict_opts) = [
+          compiler: Emily.Compiler,
+          native: true,
+          native_fallback: :raise,
+          native_compiled: true
+        ]
+
+        unquote(body)
+      end
+    end
+  end
+
   @doc """
   Assert that every element of `left` agrees with `right` within
   `atol + rtol * |right|`.
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,

From 1af35558225675a97e3ba62ba64ae901bcf6e345 Mon Sep 17 00:00:00 2001
From: ausimian <nick@ausimian.net>
Date: Sat, 6 Jun 2026 13:25:49 +1000
Subject: [PATCH 2/4] test: add native compiler lanes to the full conformance
 suites
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 test/emily/conformance/vit_full_test.exs     |  5 +-
 test/emily/conformance/whisper_full_test.exs |  4 +-
 test/support/conformance_helper.ex           | 79 +++++++++++++-------
 3 files changed, 59 insertions(+), 29 deletions(-)

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/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/support/conformance_helper.ex b/test/support/conformance_helper.ex
index 0e1e8d2..4eeeddd 100644
--- a/test/support/conformance_helper.ex
+++ b/test/support/conformance_helper.ex
@@ -30,7 +30,7 @@ defmodule Emily.ConformanceHelper do
   defmacro __using__(_opts) do
     quote do
       import Emily.ConformanceHelper,
-        only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2]
+        only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2, mode_test: 3]
 
       setup do
         Nx.default_backend(Emily.Backend)
@@ -86,36 +86,65 @@ defmodule Emily.ConformanceHelper do
   `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.
   """
-  defmacro mode_test(name, do: body) do
-    quote do
-      test unquote(name) do
-        var!(predict_opts) = []
-        unquote(body)
-      end
+  defmacro mode_test(name, opts \\ [], do: body) do
+    tag_lanes? = Keyword.get(opts, :lane_tags, true)
+
+    lanes = [
+      lane(false, name, "", [], body),
+      lane(
+        tag_lanes? && :native,
+        name,
+        " [native]",
+        [compiler: Emily.Compiler, native: true, native_fallback: :raise],
+        body
+      ),
+      lane(
+        tag_lanes? && :native_compiled,
+        name,
+        " [native_compiled]",
+        [compiler: Emily.Compiler, native: true, native_fallback: :raise, native_compiled: true],
+        body
+      )
+    ]
 
-      @tag :native
-      test unquote(name) <> " [native]" do
-        var!(predict_opts) = [
-          compiler: Emily.Compiler,
-          native: true,
-          native_fallback: :raise
-        ]
+    quote do
+      (unquote_splicing(lanes))
+    end
+  end
 
-        unquote(body)
+  # Build one `mode_test` lane: a `test` that binds `predict_opts` for the
+  # body, optionally preceded by `@tag tag`. `tag` is `false` to emit no
+  # lane tag (the `*_full` suites rely on their own `:*_full` moduletag).
+  defp lane(tag, name, suffix, predict_opts, body) do
+    name_ast =
+      if suffix == "", do: name, else: quote(do: unquote(name) <> unquote(suffix))
+
+    test =
+      quote do
+        test unquote(name_ast) do
+          var!(predict_opts) = unquote(predict_opts)
+          unquote(body)
+        end
       end
 
-      @tag :native_compiled
-      test unquote(name) <> " [native_compiled]" do
-        var!(predict_opts) = [
-          compiler: Emily.Compiler,
-          native: true,
-          native_fallback: :raise,
-          native_compiled: true
-        ]
-
-        unquote(body)
+    if tag do
+      quote do
+        @tag unquote(tag)
+        unquote(test)
       end
+    else
+      test
     end
   end
 

From ae52d9689bf039785f3c6d1aacf97335836cbd39 Mon Sep 17 00:00:00 2001
From: ausimian <nick@ausimian.net>
Date: Sat, 6 Jun 2026 13:41:57 +1000
Subject: [PATCH 3/4] feat: lower take_along_axis in the native Expr compiler
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`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`.
---
 RELEASE.md                                 |  8 +++++
 c_src/emily/opcodes.hpp                    | 14 +++++++-
 lib/emily/ir.ex                            | 15 +++++++-
 test/emily/compiler_equivalence_test.exs   | 20 +++++++++++
 test/emily/conformance/distilbert_test.exs | 16 ++++++---
 test/support/conformance_helper.ex         | 42 ++++++++++++----------
 6 files changed, 90 insertions(+), 25 deletions(-)

diff --git a/RELEASE.md b/RELEASE.md
index 67e309e..49db2e3 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -50,6 +50,14 @@
   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`.
+
 - **`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`
diff --git a/c_src/emily/opcodes.hpp b/c_src/emily/opcodes.hpp
index a4e758b..c40cdaa 100644
--- a/c_src/emily/opcodes.hpp
+++ b/c_src/emily/opcodes.hpp
@@ -153,9 +153,12 @@ 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,
 };
 
-inline constexpr int64_t kOpcodeCount = 81;
+inline constexpr int64_t kOpcodeCount = 82;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -503,6 +506,15 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &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::Concatenate: {
     if (in.empty()) {
       throw std::invalid_argument("concatenate expects >= 1 operand");
diff --git a/lib/emily/ir.ex b/lib/emily/ir.ex
index 34b5b34..ca8ae2e 100644
--- a/lib/emily/ir.ex
+++ b/lib/emily/ir.ex
@@ -132,7 +132,10 @@ 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
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -915,6 +918,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
diff --git a/test/emily/compiler_equivalence_test.exs b/test/emily/compiler_equivalence_test.exs
index 45a7467..3b3ccee 100644
--- a/test/emily/compiler_equivalence_test.exs
+++ b/test/emily/compiler_equivalence_test.exs
@@ -368,6 +368,26 @@ 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 "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 b590941..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, mode_test: 2]
+  import Emily.ConformanceHelper,
+    only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2, mode_test: 3]
 
   alias Emily.Bumblebee.FastKernels
 
@@ -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/support/conformance_helper.ex b/test/support/conformance_helper.ex
index 4eeeddd..51ddbad 100644
--- a/test/support/conformance_helper.ex
+++ b/test/support/conformance_helper.ex
@@ -96,21 +96,27 @@ defmodule Emily.ConformanceHelper do
       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(false, name, "", [], body),
+      lane([extra_tag], name, "", [], body),
       lane(
-        tag_lanes? && :native,
+        [extra_tag, tag_lanes? && :native],
         name,
         " [native]",
         [compiler: Emily.Compiler, native: true, native_fallback: :raise],
         body
       ),
       lane(
-        tag_lanes? && :native_compiled,
+        [extra_tag, tag_lanes? && :native_compiled],
         name,
         " [native_compiled]",
         [compiler: Emily.Compiler, native: true, native_fallback: :raise, native_compiled: true],
@@ -124,27 +130,25 @@ defmodule Emily.ConformanceHelper do
   end
 
   # Build one `mode_test` lane: a `test` that binds `predict_opts` for the
-  # body, optionally preceded by `@tag tag`. `tag` is `false` to emit no
-  # lane tag (the `*_full` suites rely on their own `:*_full` moduletag).
-  defp lane(tag, name, suffix, predict_opts, body) do
+  # 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))
 
-    test =
-      quote do
-        test unquote(name_ast) do
-          var!(predict_opts) = unquote(predict_opts)
-          unquote(body)
-        end
-      end
+    tag_attrs = for t <- tags, do: quote(do: @tag(unquote(t)))
+
+    quote do
+      (unquote_splicing(tag_attrs))
 
-    if tag do
-      quote do
-        @tag unquote(tag)
-        unquote(test)
+      test unquote(name_ast) do
+        var!(predict_opts) = unquote(predict_opts)
+        unquote(body)
       end
-    else
-      test
     end
   end
 

From 4578ceb5401e0fb5f414d99a1279c01505b4cf17 Mon Sep 17 00:00:00 2001
From: ausimian <nick@ausimian.net>
Date: Sat, 6 Jun 2026 13:46:18 +1000
Subject: [PATCH 4/4] test: guard opcode lockstep between Emily.IR and the C++
 enum
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 lib/emily/ir.ex                   |  8 ++++
 test/emily/opcode_parity_test.exs | 67 +++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)
 create mode 100644 test/emily/opcode_parity_test.exs

diff --git a/lib/emily/ir.ex b/lib/emily/ir.ex
index ca8ae2e..b819bb4 100644
--- a/lib/emily/ir.ex
+++ b/lib/emily/ir.ex
@@ -166,6 +166,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)
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