Merge pull request #176 from ausimian/feat/expr-compiler-conformance-native-lanes

Native + fusion compiler lanes for the conformance suites

Author: ausimian

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}
 

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),

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),

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())

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/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}
 

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/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}