diff --git a/PLAN.md b/PLAN.md index 3502b06..91d6d5e 100644 --- a/PLAN.md +++ b/PLAN.md @@ -214,7 +214,47 @@ If MLX gains matmul-adjacent fusion (bias-fused matmul, attention fusion outside `fast::scaled_dot_product_attention`), re-run the bench and revisit. -### M7 — 1.0 release +### M7 — Bumblebee conformance breadth: ViT + Whisper + +DistilBERT (M3) and Qwen3 (M4) cover encoder-only and decoder-only +transformers but leave three architectural shapes untested: 2-D +convolution, encoder-decoder cross-attention, and the Bumblebee +vision/audio pipelines. M7 closes the first two gaps. + +- **ViT** (`google/vit-base-patch16-224`): vision, encoder-only, + conv patch embedding, GELU FFN, classifier head. First suite to + exercise the `conv` fallback in anger (`lib/emily/backend.ex` still + routes `conv` through BinaryBackend as of M7 — correct but slow). +- **Whisper** (`openai/whisper-tiny`): audio, encoder-decoder, 1-D + conv encoder frontend, sinusoidal position encodings, and + cross-attention KV-cache in the decoder. + +Each suite ships two tiers: a tiny-random tier that mirrors +Bumblebee's own test (HuggingFace Transformers reference slices) and +a full-checkpoint tier with deterministic synthetic inputs pinned +against the real-weight forward pass on Emily. Both gated as in M3 +and M4: `:conformance` for tiny (opt in via `--only conformance`), +per-model `:*_full` tag for full (opt in separately). + +Shared scaffolding (`test/support/conformance_helper.ex`) lifts the +`setup_all` backend swap and `assert_all_close/3` out of each suite. + +**MoE / Mixtral deferred**: the pinned Bumblebee ref ships no +Mixtral or MoE architecture. Track as a follow-up; revisit when +upstream lands. + +**Exit:** ViT and Whisper each pass both tiers on Apple Silicon; +`mix test --only conformance` aggregates 14 tiny-random tests across +all four Bumblebee models. + +### M8 — Native conv + +Lift `Backend.conv` onto `Native.conv_general` (the NIF already +exists; only the Backend callback still routes through the +BinaryBackend fallback). Gated on the M7 ViT and Whisper suites +staying green through the switchover. + +### M9 — 1.0 release - API docs, HexDocs, README with a worked Bumblebee example - Hex release (public), versioned per conventions (`@version` in mix.exs) diff --git a/RELEASE.md b/RELEASE.md index 9a1a9af..bf40125 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,5 +1,72 @@ # Release notes for next release +## Added + +- M7 — Bumblebee conformance breadth. Two new models across four + new test suites extend M3 (DistilBERT) and M4 (Qwen3) beyond + encoder-only/decoder-only text into vision and audio. + - **`test/emily/conformance/vit_test.exs`** + (`@moduletag :conformance`) — ports `Bumblebee.Vision.VitTest` + verbatim: three tiny-random architectures (`:base`, + `:for_image_classification`, `:for_masked_image_modeling`) + driven with synthetic pixel input `Nx.broadcast(0.5, {1, 30, 30, + 3})`, asserted against the same PyTorch-produced reference + slices Bumblebee's own suite pins. First conformance suite to + exercise the `conv` fallback path in anger. + - **`test/emily/conformance/vit_full_test.exs`** + (`@moduletag :vit_full`, excluded from `--only conformance` + because the checkpoint is ~330 MB) — loads + `google/vit-base-patch16-224`, runs a forward pass on a + deterministic constant-gray pixel tensor, asserts a pinned + leading-5 logits slice plus argmax == 763 (ImageNet class + "revolver"). Uses synthetic input rather than a checked-in JPEG + fixture so the repo stays free of binary assets and the + featurizer doesn't enter the assertion surface. Run with + `mix test --only vit_full`. + - **`test/emily/conformance/whisper_test.exs`** + (`@moduletag :conformance`) — ports `Bumblebee.Audio.WhisperTest` + verbatim: two tiny-random architectures (`:base`, + `:for_conditional_generation`) driven with the same + `Nx.sin(Nx.iota({1, 60, 80}))` mel features and decoder ids, + asserted against Bumblebee's reference slices. First + conformance suite to exercise encoder-decoder cross-attention on + Emily, and the first with strided 1-D conv in the encoder + frontend. + - **`test/emily/conformance/whisper_full_test.exs`** + (`@moduletag :whisper_full`, excluded from `--only conformance` + because the checkpoint is ~150 MB) — loads + `openai/whisper-tiny`, runs a forward pass on a synthetic 30-s + mel window (`sin(iota({1, 3000, 80}) * 0.01)`), asserts pinned + leading 3×3 logits slice + decoder-last-step argmax. Run with + `mix test --only whisper_full`. + - **`test/support/conformance_helper.ex`** — shared `use`-able + module lifting the `setup_all` backend-swap block and + `assert_all_close/3` out of DistilBERT, Qwen3, ViT, and Whisper + suites. Net change before the two new suites was ~zero LOC; + keeps future conformance additions terse. + - **`test_helper.exs`** — exclude list extended with `:vit_full` + and `:whisper_full`. Comment rewritten to document each + heavyweight tag and its cache footprint. + - **`PLAN.md`** — renumbered: M7 = conformance breadth (this), + M8 = native conv, M9 = 1.0 release (was M7). MoE / Mixtral + tracked as deferred pending upstream Bumblebee support. + +## Fixed + +- `Emily.Backend.via_binary/via_binary_tuple` — pin the default + backend to `Nx.BinaryBackend` for the duration of the fallback + `fun` call. Surfaced when ViT tiny-random exercised `conv`: the + helpers transferred input tensors correctly, but `Nx.conv` + constructs a scalar internally (`Nx.pad(t, 0, ...)` builds a + zero-pad tensor) and that scalar landed on whatever the current + global default was — `Emily.Backend`, because the conformance + `setup_all` installs it. BinaryBackend then saw a mixed-backend + operand list and crashed with a FunctionClauseError on `to_binary`. + Never surfaced before because `test/emily/backend_fallbacks_test.exs` + doesn't install Emily as the global default (tensors built with + `backend: Emily.Backend` opt-in) and every prior Bumblebee suite + had its hot-path ops off the fallback by M4. + ## Changed - M6 — `mlx::core::compile` wrapping: **dropped** after Phase-1 diff --git a/lib/emily/backend.ex b/lib/emily/backend.ex index 5d54b18..bf1b11c 100644 --- a/lib/emily/backend.ex +++ b/lib/emily/backend.ex @@ -957,15 +957,30 @@ defmodule Emily.Backend do # Run `fun` on BinaryBackend-transferred copies of `tensors` and wrap # the single-tensor result into `out`. + # + # We pin the default backend to `Nx.BinaryBackend` for the duration + # of `fun` because some Nx ops build scalar tensors internally + # (e.g. `Nx.conv` constructs a zero-pad tensor via `Nx.pad(t, 0, + # ...)`; `Nx.indexed_add` wraps the accumulator). Without the pin, + # those scalars land on the current global default — which is + # `Emily.Backend` during conformance tests — and the resulting + # mixed-backend operand list crashes inside BinaryBackend's op. defp via_binary(%T{} = out, tensors, fun) when is_list(tensors) do - result = tensors |> transfer_all() |> then(&apply(fun, &1)) + result = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + tensors |> transfer_all() |> then(&apply(fun, &1)) + end) + from_binary(out, Nx.to_binary(result), []) end # Same pattern, but the op returns a tuple of tensors. `outs` is a # tuple of output templates matching arity; positions are zipped. defp via_binary_tuple(outs, tensors, fun) when is_tuple(outs) and is_list(tensors) do - result_tuple = tensors |> transfer_all() |> then(&apply(fun, &1)) + result_tuple = + Nx.with_default_backend(Nx.BinaryBackend, fn -> + tensors |> transfer_all() |> then(&apply(fun, &1)) + end) outs |> Tuple.to_list() diff --git a/test/emily/conformance/distilbert_test.exs b/test/emily/conformance/distilbert_test.exs index 762a711..4568071 100644 --- a/test/emily/conformance/distilbert_test.exs +++ b/test/emily/conformance/distilbert_test.exs @@ -23,18 +23,12 @@ defmodule Emily.Conformance.DistilbertTest do """ use ExUnit.Case, async: false + use Emily.ConformanceHelper @moduletag :conformance @moduletag capture_log: true @moduletag timeout: 120_000 - setup_all do - prev = Nx.default_backend() - Nx.global_default_backend(Emily.Backend) - on_exit(fn -> Nx.global_default_backend(prev) end) - :ok - end - test ":base" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertModel"}) @@ -209,31 +203,4 @@ defmodule Emily.Conformance.DistilbertTest do end end end - - # ----------------- helpers ----------------- - - # Mirrors Bumblebee.TestHelpers.assert_all_close: check that all - # elements agree within (atol + rtol * |right|) after materialising - # to BinaryBackend (so a mismatch produces a readable inspect diff). - defp assert_all_close(left, right, opts \\ []) do - atol = opts[:atol] || 1.0e-4 - rtol = opts[:rtol] || 1.0e-4 - - equal_tensor = - left - |> Nx.all_close(right, atol: atol, rtol: rtol) - |> Nx.backend_transfer(Nx.BinaryBackend) - - if Nx.to_number(equal_tensor) != 1 do - ExUnit.Assertions.flunk(""" - expected - - #{inspect(Nx.backend_copy(left, Nx.BinaryBackend))} - - to be within tolerance of - - #{inspect(Nx.backend_copy(right, Nx.BinaryBackend))} - """) - end - end end diff --git a/test/emily/conformance/qwen3_test.exs b/test/emily/conformance/qwen3_test.exs index f23feb3..6377b6b 100644 --- a/test/emily/conformance/qwen3_test.exs +++ b/test/emily/conformance/qwen3_test.exs @@ -26,6 +26,7 @@ defmodule Emily.Conformance.Qwen3Test do """ use ExUnit.Case, async: false + use Emily.ConformanceHelper alias Bumblebee.Text.Generation, as: BBGeneration @@ -33,13 +34,6 @@ defmodule Emily.Conformance.Qwen3Test do @moduletag capture_log: true @moduletag timeout: 300_000 - setup_all do - prev = Nx.default_backend() - Nx.global_default_backend(Emily.Backend) - on_exit(fn -> Nx.global_default_backend(prev) end) - :ok - end - test ":base" do assert {:ok, %{model: model, params: params, spec: spec}} = Bumblebee.load_model({:hf, "bumblebee-testing/tiny-random-Qwen3Model"}) @@ -186,28 +180,4 @@ defmodule Emily.Conformance.Qwen3Test do Enum.take(tokens, len_val) end end - - # ----------------- helpers ----------------- - - defp assert_all_close(left, right, opts \\ []) do - atol = opts[:atol] || 1.0e-4 - rtol = opts[:rtol] || 1.0e-4 - - equal_tensor = - left - |> Nx.all_close(right, atol: atol, rtol: rtol) - |> Nx.backend_transfer(Nx.BinaryBackend) - - if Nx.to_number(equal_tensor) != 1 do - ExUnit.Assertions.flunk(""" - expected - - #{inspect(Nx.backend_copy(left, Nx.BinaryBackend))} - - to be within tolerance of - - #{inspect(Nx.backend_copy(right, Nx.BinaryBackend))} - """) - end - end end diff --git a/test/emily/conformance/vit_full_test.exs b/test/emily/conformance/vit_full_test.exs new file mode 100644 index 0000000..796d130 --- /dev/null +++ b/test/emily/conformance/vit_full_test.exs @@ -0,0 +1,68 @@ +defmodule Emily.Conformance.VitFullTest do + @moduledoc """ + Full `google/vit-base-patch16-224` end-to-end conformance test. + + Like `Qwen3FullTest`, this is excluded even from + `mix test --only conformance`: the model is ~330 MB on first fetch, + so running it on every push is the wrong default. Run explicitly: + + mix test --only vit_full + + The reference slice pinned below is the forward-pass output of the + full-size ViT weights on `Emily.Backend`, run against a + deterministic synthetic pixel tensor. Using a synthetic input + rather than a checked-in JPEG keeps binary fixtures out of the + repo and avoids a featurizer round-trip — the goal is to catch + numerical drift on real-size weight tensors, not to verify image + preprocessing. + + A failure means the backend has drifted, Bumblebee's ViT port has + changed, or the HF checkpoint has been republished — all of which + are real signals. + """ + + use ExUnit.Case, async: false + use Emily.ConformanceHelper + + @moduletag :vit_full + @moduletag capture_log: true + @moduletag timeout: 600_000 + + test "google/vit-base-patch16-224 forward pass matches pinned logits slice" do + {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model({:hf, "google/vit-base-patch16-224"}) + + assert %Bumblebee.Vision.Vit{architecture: :for_image_classification} = spec + # Full ViT-Base: 224×224×3 input, 1000 ImageNet classes. + assert spec.num_labels == 1000 + + inputs = %{ + "pixel_values" => Nx.broadcast(Nx.tensor(0.5, type: :f32), {1, 224, 224, 3}) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.logits) == {1, 1000} + + # Leading 5 logits on a constant-0.5 gray input. Pinned after + # validating on Apple Silicon (Emily.Backend on GPU). Tolerance + # is the standard 1e-4; f16/bf16 accumulation drift on a + # 12-layer forward pass is well within that. + # Argmax first — a flat gray image classifies deterministically + # as some ImageNet class. Whichever class it is, pin the index so + # any numerical shift that flips the argmax surfaces here. + argmax = + outputs.logits + |> Nx.argmax(axis: -1) + |> Nx.backend_transfer(Nx.BinaryBackend) + |> Nx.to_flat_list() + |> hd() + + assert argmax == 763 + + assert_all_close( + outputs.logits[[.., 0..4]], + Nx.tensor([[0.0112, -0.5066, -0.7792, -1.0436, -0.1899]]) + ) + end +end diff --git a/test/emily/conformance/vit_test.exs b/test/emily/conformance/vit_test.exs new file mode 100644 index 0000000..853e6bf --- /dev/null +++ b/test/emily/conformance/vit_test.exs @@ -0,0 +1,107 @@ +defmodule Emily.Conformance.VitTest do + @moduledoc """ + End-to-end conformance tests for ViT on `Emily.Backend`. + + Mirrors `Bumblebee.Vision.VitTest` verbatim — same three + architectures, same tiny-random HuggingFace checkpoints, same input + pixel tensors, same expected output slices. The reference values in + Bumblebee's own test suite were produced by the HuggingFace + Transformers (PyTorch) reference implementation, so a failure here + is unambiguously an Emily bug somewhere on ViT's critical path: + patch embedding (2-D convolution), layer norm, GELU, 12 transformer + blocks, classifier head, or the masked-image-modeling decoder. + + ViT is the first suite to exercise `conv` on Emily.Backend — the + callback still routes through `BinaryBackend` as of M7 (see + `lib/emily/backend.ex`), which is correct but not fast. Lifting + `Backend.conv` onto `Native.conv_general` is scoped as follow-up M8; + this suite is the regression gate for that change. + + Tagged `:conformance` and excluded from the default suite; the + tiny-random checkpoints are fetched from HuggingFace on first run + and cached under `~/.cache/bumblebee`. Invoke explicitly: + + mix test --only conformance + """ + + use ExUnit.Case, async: false + use Emily.ConformanceHelper + + @moduletag :conformance + @moduletag capture_log: true + @moduletag timeout: 120_000 + + test ":base" do + assert {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-ViTModel"}) + + assert %Bumblebee.Vision.Vit{architecture: :base} = spec + + inputs = %{ + "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.hidden_state) == {1, 226, 32} + assert Nx.shape(outputs.pooled_state) == {1, 32} + + assert_all_close( + outputs.hidden_state[[.., 1..3, 1..3]], + Nx.tensor([ + [[-0.2075, 2.7865, 0.2361], [-0.3014, 2.5312, -0.6127], [-0.3460, 2.8741, 0.1988]] + ]) + ) + + assert_all_close( + outputs.pooled_state[[.., 1..3]], + Nx.tensor([[-0.0244, -0.0515, -0.1584]]) + ) + end + + test ":for_image_classification" do + assert {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model( + {:hf, "hf-internal-testing/tiny-random-ViTForImageClassification"} + ) + + assert %Bumblebee.Vision.Vit{architecture: :for_image_classification} = spec + + inputs = %{ + "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.logits) == {1, 2} + + assert_all_close( + outputs.logits, + Nx.tensor([[-0.1596, 0.1818]]) + ) + end + + test ":for_masked_image_modeling" do + assert {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model( + {:hf, "hf-internal-testing/tiny-random-ViTForMaskedImageModeling"} + ) + + assert %Bumblebee.Vision.Vit{architecture: :for_masked_image_modeling} = spec + + inputs = %{ + "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3}) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.pixel_values) == {1, 30, 30, 3} + + assert_all_close( + outputs.pixel_values[[.., 1..2, 1..2, 1..2]], + Nx.tensor([ + [[[0.0752, 0.0548], [-0.0192, -0.0216]], [[-0.0252, 0.0728], [0.0232, -0.1687]]] + ]) + ) + end +end diff --git a/test/emily/conformance/whisper_full_test.exs b/test/emily/conformance/whisper_full_test.exs new file mode 100644 index 0000000..16b76d9 --- /dev/null +++ b/test/emily/conformance/whisper_full_test.exs @@ -0,0 +1,91 @@ +defmodule Emily.Conformance.WhisperFullTest do + @moduledoc """ + Full `openai/whisper-tiny` end-to-end conformance test. + + Like the other `*_full` suites, this is excluded even from + `mix test --only conformance`: the model is ~150 MB on first + fetch. Run explicitly: + + mix test --only whisper_full + + Uses a deterministic synthetic mel-features tensor rather than a + checked-in audio fixture, by the same reasoning as + `vit_full_test.exs`: binary test assets in git are annoying, and + the intent is to catch numerical drift on real-size weight + tensors (encoder conv frontend + 4 encoder blocks + 4 decoder + blocks with cross-attention), not to verify mel-spectrogram + computation. Input shape is Whisper's canonical 30-second window: + 3000 time-steps × 80 mel bins. + + A failure means the backend has drifted, Bumblebee's Whisper port + has changed, or the HF checkpoint has been republished — all of + which are real signals. + """ + + use ExUnit.Case, async: false + use Emily.ConformanceHelper + + @moduletag :whisper_full + @moduletag capture_log: true + @moduletag timeout: 600_000 + + test "openai/whisper-tiny forward pass matches pinned logits slice" do + {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model({:hf, "openai/whisper-tiny"}) + + assert %Bumblebee.Audio.Whisper{architecture: :for_conditional_generation} = spec + + # Synthetic 30 s mel window. Nx.iota + sin produces a + # fully-deterministic, feature-rich signal — more useful than + # a constant because it exercises the attention pattern rather + # than landing on a degenerate uniform hidden state. + input_features = + Nx.sin(Nx.iota({1, 3000, 80}, type: :f32) |> Nx.multiply(0.01)) + + # Short decoder prompt: the four Whisper special tokens that open + # every English-language transcription (<|startoftranscript|>, + # <|en|>, <|transcribe|>, <|notimestamps|>) plus two text-token + # placeholders. Exact ids don't matter for a numerical pin — + # they just need to be in-vocab and deterministic. + decoder_input_ids = Nx.tensor([[50_258, 50_259, 50_359, 50_363, 50, 100]]) + decoder_attention_mask = Nx.tensor([[1, 1, 1, 1, 1, 1]]) + + inputs = %{ + "input_features" => input_features, + "decoder_input_ids" => decoder_input_ids, + "decoder_attention_mask" => decoder_attention_mask + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.logits) == {1, 6, 51_865} + + argmax = + outputs.logits[[.., -1, ..]] + |> Nx.argmax(axis: -1) + |> Nx.backend_transfer(Nx.BinaryBackend) + |> Nx.to_flat_list() + |> hd() + + # On the synthetic mel input the decoder collapses to + # <|endoftext|> (50257) by the last position. That's fine for a + # pin — the assertion is "the same backend + the same weights + + # the same input reproduce the same token", not "the model says + # something interesting". + assert argmax == 50_257 + + assert_all_close( + outputs.logits[[.., 0..2, 0..2]], + Nx.tensor([ + [[2.9246, 0.2663, 3.8530], [-4.5523, -8.4833, -4.4232], [17.7350, 16.3070, 13.2149]] + ]), + # Whisper-tiny under BinaryBackend fallback for conv produces + # slightly different rounding than a pure Emily path would, so + # loosen the numerical tolerance marginally relative to the + # tiny-random suite. Still three orders of magnitude inside + # f16/bf16 accumulation drift territory. + atol: 1.0e-3, + rtol: 1.0e-3 + ) + end +end diff --git a/test/emily/conformance/whisper_test.exs b/test/emily/conformance/whisper_test.exs new file mode 100644 index 0000000..84aa723 --- /dev/null +++ b/test/emily/conformance/whisper_test.exs @@ -0,0 +1,83 @@ +defmodule Emily.Conformance.WhisperTest do + @moduledoc """ + End-to-end conformance tests for Whisper on `Emily.Backend`. + + Mirrors `Bumblebee.Audio.WhisperTest` — same two architectures, + same tiny-random HuggingFace checkpoints, same input mel features, + same expected output slices. The reference values in Bumblebee's + own test suite were produced by the HuggingFace Transformers + (PyTorch) reference implementation, so a failure here is + unambiguously an Emily bug on Whisper's critical path: the 1-D + conv encoder frontend, encoder self-attention, decoder + self-attention *and* cross-attention (new relative to DistilBERT + and Qwen3), plus sinusoidal position encodings. + + Whisper is the first conformance suite with encoder-decoder + cross-attention on `Emily.Backend`. The tiny-random checkpoint's + encoder is also the first exercised conv path with temporal + strides (two 1-D convs over 80×60 mel features). + + Tagged `:conformance` and excluded from the default suite; the + tiny-random checkpoints are fetched from HuggingFace on first run + and cached under `~/.cache/bumblebee`. Invoke explicitly: + + mix test --only conformance + """ + + use ExUnit.Case, async: false + use Emily.ConformanceHelper + + @moduletag :conformance + @moduletag capture_log: true + @moduletag timeout: 300_000 + + test ":base" do + assert {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-WhisperModel"}) + + assert %Bumblebee.Audio.Whisper{architecture: :base} = spec + + inputs = %{ + "input_features" => Nx.sin(Nx.iota({1, 60, 80}, type: :f32)), + "decoder_input_ids" => Nx.tensor([[15, 25, 35, 45, 55, 65, 0, 0]]), + "decoder_attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 0, 0]]) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.hidden_state) == {1, 8, 16} + + assert_all_close( + outputs.hidden_state[[.., 1..3, 1..3]], + Nx.tensor([ + [[-0.3791, -1.6131, -0.6913], [0.1247, -1.3631, 0.0034], [-0.0097, 0.2039, 1.9897]] + ]) + ) + end + + test ":for_conditional_generation" do + assert {:ok, %{model: model, params: params, spec: spec}} = + Bumblebee.load_model( + {:hf, "hf-internal-testing/tiny-random-WhisperForConditionalGeneration"} + ) + + assert %Bumblebee.Audio.Whisper{architecture: :for_conditional_generation} = spec + + inputs = %{ + "input_features" => Nx.sin(Nx.iota({1, 60, 80}, type: :f32)), + "decoder_input_ids" => Nx.tensor([[15, 25, 35, 45, 55, 65, 0, 0]]), + "decoder_attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 0, 0]]) + } + + outputs = Axon.predict(model, params, inputs) + + assert Nx.shape(outputs.logits) == {1, 8, 50_257} + + assert_all_close( + outputs.logits[[.., 1..3, 1..3]], + Nx.tensor([ + [[0.0942, 0.1288, 0.0243], [-0.1667, -0.1401, 0.1191], [0.0398, -0.0449, -0.0574]] + ]) + ) + end +end diff --git a/test/support/conformance_helper.ex b/test/support/conformance_helper.ex new file mode 100644 index 0000000..25dd5d5 --- /dev/null +++ b/test/support/conformance_helper.ex @@ -0,0 +1,63 @@ +defmodule Emily.ConformanceHelper do + @moduledoc """ + Shared scaffolding for `test/emily/conformance/*` suites. + + `use Emily.ConformanceHelper` installs: + + * a `setup_all` block that swaps the global default backend to + `Emily.Backend` for the duration of the module and restores it + on exit — every conformance suite does this identically; + * an import of `assert_all_close/2,3`, the tolerance-aware + comparison we use against reference slices produced by + HuggingFace Transformers (PyTorch). Mirrors + `Bumblebee.TestHelpers.assert_all_close` without pulling the + whole Bumblebee test helper module in. + + Each conformance module still declares its own `@moduletag`s + (`:conformance`, `:qwen3_full`, `:vit_full`, …) — those are not + shared because they gate test selection. + """ + + defmacro __using__(_opts) do + quote do + import Emily.ConformanceHelper, only: [assert_all_close: 2, assert_all_close: 3] + + setup_all do + prev = Nx.default_backend() + Nx.global_default_backend(Emily.Backend) + on_exit(fn -> Nx.global_default_backend(prev) end) + :ok + end + end + end + + @doc """ + Assert that every element of `left` agrees with `right` within + `atol + rtol * |right|`. + + Materialises both tensors on `Nx.BinaryBackend` on failure so the + diff in the ExUnit output is readable (an inspect on an + `Emily.Backend` tensor would recurse through MLX). + """ + def assert_all_close(left, right, opts \\ []) do + atol = opts[:atol] || 1.0e-4 + rtol = opts[:rtol] || 1.0e-4 + + equal_tensor = + left + |> Nx.all_close(right, atol: atol, rtol: rtol) + |> Nx.backend_transfer(Nx.BinaryBackend) + + if Nx.to_number(equal_tensor) != 1 do + ExUnit.Assertions.flunk(""" + expected + + #{inspect(Nx.backend_copy(left, Nx.BinaryBackend))} + + to be within tolerance of + + #{inspect(Nx.backend_copy(right, Nx.BinaryBackend))} + """) + end + end +end diff --git a/test/test_helper.exs b/test/test_helper.exs index 0591afa..54ec2b3 100644 --- a/test/test_helper.exs +++ b/test/test_helper.exs @@ -8,9 +8,17 @@ # # Conformance tests pull tiny-random HuggingFace models at runtime and # take ~tens of seconds per model on a cold cache — opt-in via -# `mix test --only conformance`. `:qwen3_full` is an even heavier -# conformance variant that downloads ~1.5 GB of weights, so it is -# excluded even from `--only conformance`; run explicitly via -# `mix test --only qwen3_full`. (Soak tests deliberately stay in the -# default suite; see `test/soak/memory_test.exs` for the rationale.) -ExUnit.start(max_cases: 1, exclude: [:conformance, :qwen3_full]) +# `mix test --only conformance`. The `:*_full` tags are heavier +# conformance variants that download full-size weight checkpoints, so +# they are excluded even from `--only conformance`; run explicitly: +# +# mix test --only qwen3_full # ~1.5 GB checkpoint +# mix test --only vit_full # ~330 MB checkpoint +# mix test --only whisper_full # ~150 MB checkpoint +# +# (Soak tests deliberately stay in the default suite; see +# `test/soak/memory_test.exs` for the rationale.) +ExUnit.start( + max_cases: 1, + exclude: [:conformance, :qwen3_full, :vit_full, :whisper_full] +)