test: add native compiler lanes to the full conformance suites

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.

Author: ausimian

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}
 

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}
 

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