Gate the native Whisper speech_to_text serving end-to-end

RELEASE.md

@@ -67,6 +67,14 @@
   `speech_to_text` serving — the log-mel featurizer's STFT — so that path
   now compiles fully native too.
 
+- **`indexed_put` / `indexed_add` (scatter) lower natively** — both now
+  compile under the native single-NIF path for MLX-scatter-compatible index
+  layouts (the same layout the native gather already requires), mirroring
+  `Emily.Backend`'s scatter — split the index tensor into per-axis s32
+  arrays, reshape updates into MLX's layout, then `mx::scatter` (overwrite)
+  / `mx::scatter_add` (accumulate) bit-for-bit. Layouts MLX can't scatter
+  still route through the evaluator under `native_fallback: :eval`.
+
 - **Window (pooling) ops lower natively — forward and backward.** The
   forward window family (`window_sum`/`window_max`/`window_min`/
   `window_product`, i.e. average and max pooling), the select-and-scatter
@@ -177,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

c_src/emily/opcodes.hpp

@@ -175,9 +175,13 @@ enum class Opcode : int64_t {
   Ifftn = 89,  // complex -> complex (inverse)
   Rfftn = 90,  // real -> complex (half spectrum)
   Irfftn = 91, // complex half-spectrum -> real
+  // Scatter (Nx.indexed_put / indexed_add). operands [target, updates,
+  // idx0, ...] (one s32 index array per scattered axis); iattrs [[axes...]].
+  Scatter = 92,    // overwrite (last write wins on duplicates)
+  ScatterAdd = 93, // accumulate
 };
 
-inline constexpr int64_t kOpcodeCount = 92;
+inline constexpr int64_t kOpcodeCount = 94;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -691,6 +695,21 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
                            emily::to_mlx_shape(attr_at(iattrs, 0, "irfftn")),
                            emily::to_int_vec(attr_at(iattrs, 1, "irfftn")),
                            mx::fft::FFTNorm::Backward, s);
+  // --- Scatter (shares the eager index.cpp entry points) ---
+  case Opcode::Scatter:
+  case Opcode::ScatterAdd: {
+    if (in.size() < 3) {
+      throw std::invalid_argument(
+          "scatter expects >= 3 operands (target, updates, >=1 index), got " +
+          std::to_string(in.size()));
+    }
+    // operands [target, updates, idx0, ...]; the index arrays follow updates.
+    std::vector<mx::array> indices(in.begin() + 2, in.end());
+    auto axes = emily::to_int_vec(attr0(iattrs, "scatter"));
+    return op == Opcode::Scatter
+               ? mx::scatter(in[0], indices, in[1], axes, s)
+               : mx::scatter_add(in[0], indices, in[1], axes, s);
+  }
   }
   throw std::invalid_argument("unknown opcode " +
                               std::to_string(static_cast<int64_t>(op)));

lib/emily/ir.ex

@@ -153,7 +153,12 @@ defmodule Emily.IR do
     fftn: 88,
     ifftn: 89,
     rfftn: 90,
-    irfftn: 91
+    irfftn: 91,
+    # Scatter (Nx.indexed_put / indexed_add). operands [target, updates,
+    # idx0, ...] (one s32 index array per scattered axis); iattrs [[axes...]].
+    # scatter overwrites (last-write on duplicates); scatter_add accumulates.
+    scatter: 92,
+    scatter_add: 93
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -661,6 +666,39 @@ defmodule Emily.IR do
     coerce(r, t.type, state)
   end
 
+  # indexed_put / indexed_add (Nx scatter). Mirrors Emily.Backend's
+  # apply_scatter: split the {..., R} index tensor into R per-axis s32 index
+  # arrays, reshape `updates` into MLX's scatter layout, then mx::scatter
+  # (overwrite) / mx::scatter_add (accumulate). Only MLX-scatter-compatible
+  # index layouts lower; others raise (no fallback) — the Evaluator handles
+  # them via_binary under `native_fallback: :eval`, matching native gather.
+  # Operands [target, updates, idx0, ...]; iattrs [[axes...]].
+  defp lower_op(
+         %T{data: %Nx.Defn.Expr{op: op, args: [target, indices, updates, opts]}} = t,
+         state
+       )
+       when op in [:indexed_put, :indexed_add] do
+    axes = opts[:axes] || Enum.to_list(0..(tuple_size(target.shape) - 1)//1)
+    indices_shape = Tuple.to_list(indices.shape)
+
+    unless scatter_gather_compatible?(indices_shape, axes) do
+      raise ArgumentError,
+            "Emily Expr compiler: #{op} index layout #{inspect(indices_shape)} for axes " <>
+              "#{inspect(axes)} is not MLX-scatter-compatible (no fallback)."
+    end
+
+    {rt, state} = lower_node(target, state)
+    {rx, state} = lower_node(indices, state)
+    {idx_refs, state} = split_indices_for_gather(rx, indices_shape, length(axes), state)
+
+    {ru, state} = lower_node(updates, state)
+    updates_shape = updates_shape_for_scatter(indices_shape, target.shape, axes)
+    {ru, state} = emit(state, :reshape, [ru], [updates_shape])
+
+    opcode = if op == :indexed_put, do: :scatter, else: :scatter_add
+    emit_coerced(state, opcode, [rt, ru | idx_refs], [axes], t.type)
+  end
+
   # put_slice(src, start_indices, slice): write `slice` into `src` at
   # `start_indices`. Mirrors Emily.Backend.put_slice/4 (cast src + update
   # to out.type), but supports RUNTIME (tensor) start indices — the decode
@@ -832,7 +870,7 @@ defmodule Emily.IR do
   # gather/index there clamps rather than faults, so the discarded value
   # never changes the result; a hard-faulting op on a not-taken path would
   # diverge from the Evaluator's lazy single-branch eval.
-  defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, last]}} = t, state) do
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, %T{} = last]}} = t, state) do
     {last_ref, state} = lower_node(last, state)
 
     {result, state} =
@@ -846,6 +884,52 @@ defmodule Emily.IR do
     coerce(result, t.type, state)
   end
 
+  # Multi-output cond: each branch returns a tuple of tensors. Lower to one
+  # `where`-chain per leaf position — identical wholesale-select semantics to
+  # the single-output case above (the predicate is a whole-tensor scalar bool;
+  # every branch is still computed). Returns a `{:multi_refs, [...]}` handle
+  # that `:elem` projects; sibling `:elem`s share it via lower_node's memo.
+  # A nested / non-tensor container raises (no fallback path for it yet).
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, last]}}, state)
+       when is_tuple(last) do
+    last_leaves = Tuple.to_list(last)
+
+    unless tensors?(last_leaves) and
+             Enum.all?(clauses, fn {_p, b} -> is_tuple(b) and tensors?(Tuple.to_list(b)) end) do
+      raise ArgumentError,
+            "Emily Expr compiler: cond over a nested / non-tensor container is not " <>
+              "lowered yet (only a flat tuple of tensors)."
+    end
+
+    # Lower each predicate (cast to pred) + its branch's leaf refs once, so the
+    # per-leaf where-chains share them.
+    {clauses, state} =
+      Enum.map_reduce(clauses, state, fn {pred, body}, st ->
+        {pred_ref, st} = lower_node(pred, st)
+        {pred_ref, st} = emit(st, :astype, [pred_ref], [[dtype_code({:pred, 1})]])
+        {body_refs, st} = Enum.map_reduce(Tuple.to_list(body), st, &lower_node/2)
+        {{pred_ref, body_refs}, st}
+      end)
+
+    rev = Enum.reverse(clauses)
+
+    {refs, state} =
+      last_leaves
+      |> Enum.with_index()
+      |> Enum.map_reduce(state, fn {leaf, j}, st ->
+        {last_ref, st} = lower_node(leaf, st)
+
+        {result, st} =
+          Enum.reduce(rev, {last_ref, st}, fn {pred_ref, body_refs}, {else_ref, st2} ->
+            emit(st2, :where, [pred_ref, Enum.at(body_refs, j), else_ref])
+          end)
+
+        coerce(result, leaf.type, st)
+      end)
+
+    {{:multi_refs, refs}, state}
+  end
+
   # attach_token: sequences a token (hooks) before `expr`. With no active
   # hook the token is a no-op, so pass through to the inner expr. Hooks
   # would need a callback into Elixir mid-graph (program-split) — deferred.
@@ -917,11 +1001,14 @@ defmodule Emily.IR do
       {{:multi, base, _arity}, state} ->
         {{:instr, base + i}, state}
 
+      {{:multi_refs, refs}, state} ->
+        {Enum.at(refs, i), state}
+
       {_handle, _state} ->
         raise ArgumentError,
               "Emily Expr compiler: :elem projects a tuple-producing op it can't " <>
-                "lower yet (only `while` produces projectable tuples today; other " <>
-                "multi-output ops are unsupported)."
+                "lower yet (only `while` and tuple `cond` produce projectable " <>
+                "tuples today; other multi-output ops are unsupported)."
     end
   end
 
@@ -1119,6 +1206,8 @@ defmodule Emily.IR do
   defp bool_int(true), do: 1
   defp bool_int(false), do: 0
 
+  defp tensors?(list), do: Enum.all?(list, &match?(%T{}, &1))
+
   defp float_like?({kind, _}) when kind in [:f, :bf, :c], do: true
   defp float_like?(_), do: false
 
@@ -1285,4 +1374,19 @@ defmodule Emily.IR do
     rank = tuple_size(input_shape)
     for i <- 0..(rank - 1)//1, do: if(i in axes_set, do: 1, else: elem(input_shape, i))
   end
+
+  # Rewrap Nx's updates shape {batch ++ non_indexed_dims} into MLX's scatter
+  # layout {batch ++ per_axis_slot}, where per_axis_slot has length
+  # rank(target) with 1 on indexed axes and target_shape[i] elsewhere. Mirrors
+  # Emily.Backend.updates_shape_for_scatter/3.
+  defp updates_shape_for_scatter(indices_shape, target_shape, axes) do
+    batch = Enum.take(indices_shape, length(indices_shape) - 1)
+    axes_set = MapSet.new(axes)
+    rank = tuple_size(target_shape)
+
+    trailing =
+      for i <- 0..(rank - 1)//1, do: if(i in axes_set, do: 1, else: elem(target_shape, i))
+
+    batch ++ trailing
+  end
 end

test/emily/compiler_control_flow_test.exs

@@ -45,6 +45,34 @@ defmodule Emily.CompilerControlFlowTest do
     acc
   end
 
+  # Tuple-returning cond (multi-output): both branches produce a {vec, vec}
+  # tuple; `:elem` projects each leaf out of the lowered per-leaf where-chains.
+  defn tuple_if_fn(x) do
+    {a, b} =
+      if Nx.greater(Nx.sum(x), 0) do
+        {Nx.multiply(x, 2), Nx.add(x, 1)}
+      else
+        {Nx.negate(x), Nx.subtract(x, 1)}
+      end
+
+    Nx.add(a, b)
+  end
+
+  # Multi-clause tuple cond with leaves of DIFFERENT shapes (vector + scalar),
+  # so each leaf's where-chain is built and coerced independently.
+  defn tuple_cond3_fn(x) do
+    s = Nx.sum(x)
+
+    {vec, scalar} =
+      cond do
+        Nx.greater(s, 10) -> {Nx.multiply(x, 10), Nx.reduce_max(x)}
+        Nx.greater(s, 0) -> {Nx.multiply(x, 2), Nx.product(x)}
+        true -> {Nx.negate(x), s}
+      end
+
+    Nx.add(vec, scalar)
+  end
+
   defp equiv(fun, x) do
     native = Nx.Defn.jit(fun, @native).(x)
     eval = Nx.Defn.jit(fun, @eval).(x)
@@ -71,6 +99,18 @@ defmodule Emily.CompilerControlFlowTest do
         equiv(&nested_if_fn/1, Nx.tensor(data, backend: Emily.Backend))
       end
     end
+
+    test "tuple-returning if (multi-output cond) projects each leaf" do
+      for data <- [[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]] do
+        equiv(&tuple_if_fn/1, Nx.tensor(data, backend: Emily.Backend))
+      end
+    end
+
+    test "multi-clause tuple cond with mixed-shape leaves matches the evaluator" do
+      for data <- [[5.0, 4.0, 3.0], [1.0, 1.0, 1.0], [-2.0, -2.0, -2.0]] do
+        equiv(&tuple_cond3_fn/1, Nx.tensor(data, backend: Emily.Backend))
+      end
+    end
   end
 
   describe "defn while" do

test/emily/compiler_equivalence_test.exs

@@ -395,6 +395,32 @@ defmodule Emily.CompilerEquivalenceTest do
     end
   end
 
+  describe "scatter (indexed_put / indexed_add)" do
+    test "indexed_put / indexed_add into a 2-D target match the evaluator" do
+      target = et([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+      # Two distinct cells in the {2,3} grid -> order-independent.
+      idx = Nx.tensor([[0, 0], [1, 2]], type: :s64, backend: Emily.Backend)
+      upd = et([10.0, 20.0])
+
+      assert_equiv(fn t, i, u -> Nx.indexed_put(t, i, u) end, [target, idx, upd])
+      assert_equiv(fn t, i, u -> Nx.indexed_add(t, i, u) end, [target, idx, upd])
+    end
+
+    test "indexed_add accumulates duplicate indices (same MLX kernel both paths)" do
+      target = et([0.0, 0.0, 0.0, 0.0])
+      idx = Nx.tensor([[1], [1], [3]], type: :s64, backend: Emily.Backend)
+      upd = et([5.0, 7.0, 2.0])
+      assert_equiv(fn t, i, u -> Nx.indexed_add(t, i, u) end, [target, idx, upd])
+    end
+
+    test "indexed_put on a partial axis set (axes: [0], whole-row writes) matches" do
+      target = et([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
+      idx = Nx.tensor([[0], [2]], type: :s64, backend: Emily.Backend)
+      upd = et([[10.0, 11.0, 12.0], [13.0, 14.0, 15.0]])
+      assert_equiv(fn t, i, u -> Nx.indexed_put(t, i, u, axes: [0]) end, [target, idx, upd])
+    end
+  end
+
   describe "window reductions (pooling forward)" do
     test "2x2 maxpool / sumpool / minpool (CNN-shaped) match the Evaluator" do
       # {batch, channels, h, w}; pool only the spatial axes, stride 2.

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