M7: Bumblebee conformance breadth — ViT + Whisper

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)

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

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

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

test/emily/conformance/qwen3_test.exs

@@ -26,20 +26,14 @@ defmodule Emily.Conformance.Qwen3Test do
   """
 
   use ExUnit.Case, async: false
+  use Emily.ConformanceHelper
 
   alias Bumblebee.Text.Generation, as: BBGeneration
 
   @moduletag :conformance
   @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

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