Gate the native Whisper speech_to_text serving end-to-end - #186
Merged
ausimian merged 3 commits intoJun 6, 2026
Conversation
Nx.indexed_put and Nx.indexed_add (scatter) now lower under
`compiler: Emily.Compiler, native: true` for MLX-scatter-compatible index
layouts, instead of forcing a graceful fallback to the evaluator. The MLX
ops and eager NIFs already existed (Native.{scatter,scatter_add} ->
mx::{scatter,scatter_add}); only the compiler path was missing.
- opcodes.hpp: add Scatter/ScatterAdd (92-93), bump kOpcodeCount to 94,
dispatch to mx::scatter (overwrite) / mx::scatter_add (accumulate).
operands [target, updates, idx0, ...]; iattrs [[axes...]].
- ir.ex: add the opcodes; lower :indexed_put/:indexed_add (one shared
clause), mirroring Emily.Backend's apply_scatter — reuse the existing
gather index-split + scatter_gather_compatible? helpers, port
updates_shape_for_scatter. Incompatible index layouts raise (no
fallback), matching the native gather; the evaluator handles them under
native_fallback: :eval.
- compiler_equivalence_test.exs: native-vs-evaluator bit-identical cases
for indexed_put/indexed_add (2-D grid, duplicate-index accumulation,
partial-axis whole-row writes).
Another op surfaced while running the native compiler on the Whisper
livebook.
A `cond`/`if` whose branches return a tuple crashed the native compiler
with a FunctionClauseError: the :cond lowering passed the tuple branch
straight to lower_node/2, which only matches a single tensor. Worse, a
FunctionClauseError escapes the compiler's graceful-fallback rescue (which
only traps ArgumentError), so it hard-faulted instead of degrading to the
evaluator. Surfaced by a Whisper encoder forward under native: true
(transpose(elem(cond -> {f32[1,1500,6,64] x N}, i))).
Lower a tuple cond to one where-chain per leaf position — identical
wholesale-select semantics to the single-output cond (the predicate is a
whole-tensor scalar bool; every branch is still computed) — and return a
{:multi_refs, [...]} handle that :elem projects, shared across sibling
:elems via lower_node's memo (like while). A nested / non-tensor container
raises a clean ArgumentError, so it falls back gracefully rather than
crashing.
- ir.ex: split the :cond clause (single-tensor head unchanged; new
is_tuple head for the multi-output case); extend :elem to project
{:multi_refs, refs}; add the tensors?/1 guard.
- compiler_control_flow_test.exs: tuple if + multi-clause tuple cond with
mixed-shape (vector + scalar) leaves, native-vs-evaluator bit-identical.
Add a :whisper_full case that runs the real openai/whisper-tiny speech_to_text serving on synthetic audio under native_fallback: :raise — a no-fallback gate over the WHOLE graph: the mel featurizer's STFT (fft), the encoder/decoder forward, and the autoregressive decode loop (the multi-output cond, indexed_put cache writes, dynamic slices). An un-lowerable op raises rather than silently degrading. The existing whisper forward mode_test never reaches these — it feeds pre-computed mel features through a single Axon.predict, exercising neither the featurizer nor the generation loop. That blind spot is exactly where this session's fft / indexed_put / tuple-cond gaps lived; this locks the composed path so it can't silently regress. Reuses the checkpoint whisper_full_test already downloads (no new fetch). Also adds the RELEASE.md "Fixed" note for the tuple-cond crash and the "Whisper now fully native end-to-end" milestone.
ausimian
changed the base branch from
feat/expr-compiler-cond-tuple
to
feat/expr-compiler
June 6, 2026 07:04
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Locks in this session's milestone: the full Whisper
speech_to_textserving compiles fully native (featurizer STFT + encoder/decoder + decode
loop). Without a gate, the FFT /
indexed_put/ tuple-condlowerings couldsilently regress and re-introduce fallbacks — and the existing Whisper tests
wouldn't catch it.
Why the existing tests don't cover this
whisper_test.exs/whisper_full_test.exsfeed pre-computed mel featuresthrough a single
Axon.predictforward. They never run the featurizer (so:fftwas untested) or the generation decode loop (where the tuplecondlived).
generation_native_test.exscovers text generation, not Whisper. Soall three of this session's gaps sat in an untested seam.
The gate
A
:whisper_fullcase runs the realopenai/whisper-tinyspeech_to_textserving on ~1 s of synthetic audio under
compiler: Emily.Compiler, native: true, native_fallback: :raise.:raisemakes it a no-fallback gate over the whole graph — the featurizer's
fft,the encoder/decoder forward, and the autoregressive decode loop (the
multi-output
cond,indexed_putcache writes, dynamic slices). Anun-lowerable op raises instead of degrading. Reuses the checkpoint
whisper_full_testalready downloads — no new fetch.The transcription itself isn't pinned (synthetic audio → arbitrary tokens, and
greedy decode can wobble); reaching the end under
:raiseis the assertion.Verified locally (cached checkpoint): 1 test, 0 failures, 2.6 s.
Why tiny-random can't host this
I tried — tiny-random Whisper's encoder maxes at 30 source positions, but the
featurizer pads to Whisper's 30 s window (1500 positions), and fractional
num_secondsbreaks the integer pad. So the end-to-end serving needs areal-size model; hence
:whisper_full(opt-in), reusing the existing download.Also here
The RELEASE.md "Fixed" note for the tuple-
condcrash (PR #185 didn't carryone) + the "Whisper now fully native end-to-end" milestone line.
Top of the stack:
…-scatter(#184) →…-cond-tuple(#185) → this. Baseretargets to
feat/expr-compileras the lower PRs merge.