Merge pull request #186 from ausimian/feat/expr-compiler-whisper-native-gate

Gate the native Whisper speech_to_text serving end-to-end

Author: ausimian

RELEASE.md

@@ -185,6 +185,18 @@
 
 ### Fixed
 
+- **A tuple-returning `cond` hard-crashed the native compiler instead of
+  lowering.** A `cond`/`if` whose branches return a tuple (multi-output) hit a
+  `FunctionClauseError` in the lowerer — and because that isn't an
+  `ArgumentError`, it escaped the graceful-fallback rescue and faulted rather
+  than degrading to the evaluator. It now lowers to one `where`-chain per leaf
+  (same wholesale-select semantics as a single-output `cond`), projected by
+  `:elem`; a nested/non-tensor container raises cleanly (graceful fallback).
+  Surfaced by a Whisper `speech_to_text` serving — with this, plus the native
+  `fft` and `indexed_put` lowering above, the **full Whisper serving
+  (featurizer STFT + encoder/decoder + autoregressive decode loop) compiles
+  fully native end-to-end**, gated by `whisper_full_test.exs`.
+
 - **Dilated window reductions (`window_dilations > 1`) returned wrong values.**
   `window_sum`/`window_max`/`window_min`/`window_product` with a dilated kernel
   silently produced garbage for windows past the first stride positions, on both

test/emily/conformance/whisper_full_test.exs

@@ -84,6 +84,44 @@ defmodule Emily.Conformance.WhisperFullTest do
     )
   end
 
+  test "speech_to_text serving lowers fully native — featurizer + decode loop, no fallback" do
+    repo = {:hf, "openai/whisper-tiny"}
+    {:ok, whisper} = Bumblebee.load_model(repo)
+    {:ok, featurizer} = Bumblebee.load_featurizer(repo)
+    {:ok, tokenizer} = Bumblebee.load_tokenizer(repo)
+    {:ok, generation_config} = Bumblebee.load_generation_config(repo)
+
+    # The gate is "does the whole graph lower", not transcription quality —
+    # cap the decode loop so it stays fast.
+    generation_config = Bumblebee.configure(generation_config, max_new_tokens: 4)
+
+    # `native_fallback: :raise` makes this a no-fallback gate over the ENTIRE
+    # serving graph — the mel featurizer's STFT (`fft`), the encoder/decoder
+    # forward, and the autoregressive decode loop (the multi-output `cond` in
+    # the encoder attention, `indexed_put` cache writes, dynamic slices). Any
+    # op the Expr compiler can't lower raises here rather than silently
+    # degrading. The `mode_test` forward pass above never reaches these: it
+    # feeds pre-computed mel features through a single `Axon.predict`, so it
+    # exercises neither the featurizer nor the generation loop.
+    serving =
+      Bumblebee.Audio.speech_to_text_whisper(whisper, featurizer, tokenizer, generation_config,
+        defn_options: [compiler: Emily.Compiler, native: true, native_fallback: :raise]
+      )
+
+    # ~1 s of deterministic synthetic audio. The featurizer pads it to
+    # Whisper's 30 s window, so the encoder still runs the full 1500-position
+    # path (the shape that surfaced the multi-output cond).
+    audio = Nx.sin(Nx.iota({16_000}, type: :f32) |> Nx.multiply(0.02))
+
+    %{chunks: chunks} = Nx.Serving.run(serving, audio)
+
+    # Reaching here is the gate: the full path lowered native with zero
+    # fallback. The output is only sanity-checked (synthetic audio decodes to
+    # arbitrary tokens; the transcription itself is not pinned).
+    assert is_list(chunks) and chunks != []
+    assert Enum.all?(chunks, &is_binary(&1.text))
+  end
+
   @tag :fast_kernels_full
   test "Whisper-tiny with fused MLX kernels matches the pinned argmax within widened tolerance" do
     {:ok, %{model: model, params: params}} =