Lower take_along_axis natively + opcode lockstep guard

RELEASE.md

@@ -50,6 +50,14 @@
   selection on the native path. Remaining gaps (`gather`/scatter,
   pooling/`window_*`, cumulative) continue to work via the graceful fallback.
 
+- **`take_along_axis` lowers natively** — `Nx.take_along_axis` (the
+  `Nx.Block.TakeAlongAxis` block) now compiles under the native single-NIF
+  path, mirroring `Emily.Backend.native_take_along_axis/4` (cast indices to
+  s32, then `mlx::core::take_along_axis`) bit-for-bit. This was the last op
+  forcing a fallback in `Bumblebee.Text.question_answering`'s answer-span
+  gather, so a DistilBERT question-answering `Nx.Serving` forward now runs
+  fully native — and fused — under `native_fallback: :raise`.
+
 - **`Bumblebee.Text.generation` compiles fully native — greedy and sampling.**
   The headline result: an end-to-end Bumblebee generation (the transformer
   forward, the `defn while` decode loop, dynamic KV-cache writes, `cumsum`

c_src/emily/opcodes.hpp

@@ -153,9 +153,12 @@ enum class Opcode : int64_t {
   Gather = 79,
   // Stack tensors along a new axis. operands [t0, t1, ...]; iattrs [[axis]]
   Stack = 80,
+  // Gather along one axis with a same-rank s32 index tensor.
+  // operands [input, indices]; iattrs [[axis]]
+  TakeAlongAxis = 81,
 };
 
-inline constexpr int64_t kOpcodeCount = 81;
+inline constexpr int64_t kOpcodeCount = 82;
 
 // Quant mode code (Emily.IR @quant_modes) -> MLX mode string.
 inline std::string qmode_from_code(int64_t code) {
@@ -503,6 +506,15 @@ inline mx::array dispatch_op(Opcode op, const std::vector<mx::array> &in,
     int axis = emily::checked_int(scalar_at(iattrs, 0, "take"), "axis");
     return mx::take(in[0], in[1], axis, s);
   }
+  case Opcode::TakeAlongAxis: {
+    if (in.size() != 2) {
+      throw std::invalid_argument("take_along_axis expects 2 operands, got " +
+                                  std::to_string(in.size()));
+    }
+    int axis =
+        emily::checked_int(scalar_at(iattrs, 0, "take_along_axis"), "axis");
+    return mx::take_along_axis(in[0], in[1], axis, s);
+  }
   case Opcode::Concatenate: {
     if (in.empty()) {
       throw std::invalid_argument("concatenate expects >= 1 operand");

lib/emily/ir.ex

@@ -132,7 +132,10 @@ defmodule Emily.IR do
     cummin: 78,
     # multi-axis gather: operands [input, idx0, ...]; iattrs [[axes],[slice_sizes]]
     gather: 79,
-    stack: 80
+    stack: 80,
+    # take_along_axis: gather along one axis with a same-rank s32 index
+    # tensor. operands [input, indices]; iattrs [[axis]].
+    take_along_axis: 81
   }
 
   # Quant mode string -> code; decoded by qmode_from_code in
@@ -163,6 +166,14 @@ defmodule Emily.IR do
           outputs: [ref()]
         }
 
+  @doc """
+  The opcode name -> wire-value map. Exposed for the opcode-parity test,
+  which checks these stay in lockstep with the `Opcode` enum and
+  `kOpcodeCount` in `c_src/emily/opcodes.hpp`.
+  """
+  @spec opcodes() :: %{atom() => non_neg_integer()}
+  def opcodes, do: @opcodes
+
   @doc "Numeric wire value for an opcode name."
   @spec opcode(atom()) :: non_neg_integer()
   def opcode(name) when is_map_key(@opcodes, name), do: Map.fetch!(@opcodes, name)
@@ -915,6 +926,16 @@ defmodule Emily.IR do
     emit_coerced(state, :take, [ri, rx], [[axis]], t.type)
   end
 
+  # Nx.take_along_axis (Nx.Block.TakeAlongAxis). Mirrors
+  # Emily.Backend.native_take_along_axis/4: cast indices to s32, then
+  # mx::take_along_axis along `axis`.
+  defp lower_block(%Nx.Block.TakeAlongAxis{axis: axis}, [input, indices], _expr, t, state) do
+    {ri, state} = lower_node(input, state)
+    {rx, state} = lower_node(indices, state)
+    {rx, state} = emit(state, :astype, [rx], [[dtype_code({:s, 32})]])
+    emit_coerced(state, :take_along_axis, [ri, rx], [[axis]], t.type)
+  end
+
   # Cumulative families. Like Emily.Backend.block/4, the last-axis case uses
   # the native MLX `cumsum`/`cumprod`/`cummax`/`cummin` kernel; interior axes
   # (which MLX can't always factor) fall back to the block's composed

test/emily/compiler_equivalence_test.exs

@@ -368,6 +368,26 @@ defmodule Emily.CompilerEquivalenceTest do
     end
   end
 
+  describe "take_along_axis" do
+    test "along the last axis matches the Evaluator" do
+      x = et([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0], [9.0, 10.0, 11.0, 12.0]])
+      idx = Nx.tensor([[3, 0], [1, 2], [0, 3]], type: :s64, backend: Emily.Backend)
+      assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 1) end, [x, idx])
+    end
+
+    test "along axis 0 matches" do
+      x = et([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+      idx = Nx.tensor([[0, 1, 0], [1, 0, 1], [0, 0, 1]], type: :s64, backend: Emily.Backend)
+      assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 0) end, [x, idx])
+    end
+
+    test "3-D gather along the last axis (transformer-shaped) matches" do
+      x = Nx.iota({1, 2, 4}, type: :f32, backend: Emily.Backend) |> Nx.divide(8.0)
+      idx = Nx.tensor([[[3, 1], [0, 2]]], type: :s32, backend: Emily.Backend)
+      assert_equiv(fn t, i -> Nx.take_along_axis(t, i, axis: 2) end, [x, idx])
+    end
+  end
+
   describe "dynamic put_slice (KV-cache write)" do
     test "put_slice at a runtime offset matches the Evaluator" do
       # {batch, n_kv_heads, max_len, head_dim} KV buffer; write one token.

test/emily/conformance/distilbert_test.exs

@@ -33,7 +33,8 @@ defmodule Emily.Conformance.DistilbertTest do
 
   use ExUnit.Case, async: false
 
-  import Emily.ConformanceHelper, only: [assert_all_close: 2, assert_all_close: 3]
+  import Emily.ConformanceHelper,
+    only: [assert_all_close: 2, assert_all_close: 3, mode_test: 2, mode_test: 3]
 
   alias Emily.Bumblebee.FastKernels
 
@@ -52,7 +53,7 @@ defmodule Emily.Conformance.DistilbertTest do
     :ok
   end
 
-  test ":base" do
+  mode_test ":base" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertModel"})
 
@@ -63,7 +64,7 @@ defmodule Emily.Conformance.DistilbertTest do
       "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.hidden_state) == {1, 10, 32}
 
@@ -75,7 +76,7 @@ defmodule Emily.Conformance.DistilbertTest do
     )
   end
 
-  test ":for_masked_language_modeling" do
+  mode_test ":for_masked_language_modeling" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-DistilBertForMaskedLM"})
 
@@ -86,7 +87,7 @@ defmodule Emily.Conformance.DistilbertTest do
       "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.logits) == {1, 10, 1124}
 
@@ -98,7 +99,7 @@ defmodule Emily.Conformance.DistilbertTest do
     )
   end
 
-  test ":for_sequence_classification" do
+  mode_test ":for_sequence_classification" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(
                {:hf, "hf-internal-testing/tiny-random-DistilBertForSequenceClassification"}
@@ -111,14 +112,14 @@ defmodule Emily.Conformance.DistilbertTest do
       "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.logits) == {1, 2}
 
     assert_all_close(outputs.logits, Nx.tensor([[-0.0047, -0.0103]]))
   end
 
-  test ":for_token_classification" do
+  mode_test ":for_token_classification" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(
                {:hf, "hf-internal-testing/tiny-random-DistilBertForTokenClassification"}
@@ -131,7 +132,7 @@ defmodule Emily.Conformance.DistilbertTest do
       "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.logits) == {1, 10, 2}
 
@@ -141,7 +142,7 @@ defmodule Emily.Conformance.DistilbertTest do
     )
   end
 
-  test ":for_question_answering" do
+  mode_test ":for_question_answering" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(
                {:hf, "hf-internal-testing/tiny-random-DistilBertForQuestionAnswering"}
@@ -154,7 +155,7 @@ defmodule Emily.Conformance.DistilbertTest do
       "attention_mask" => Nx.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.start_logits) == {1, 10}
     assert Nx.shape(outputs.end_logits) == {1, 10}
@@ -170,7 +171,7 @@ defmodule Emily.Conformance.DistilbertTest do
     )
   end
 
-  test ":for_multiple_choice" do
+  mode_test ":for_multiple_choice" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(
                {:hf, "hf-internal-testing/tiny-random-DistilBertForMultipleChoice"}
@@ -183,7 +184,7 @@ defmodule Emily.Conformance.DistilbertTest do
       "attention_mask" => Nx.tensor([[[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]]])
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.logits) == {1, 1}
 
@@ -228,15 +229,22 @@ defmodule Emily.Conformance.DistilbertTest do
     # (vocab 30522) with a tiny-random model (1124-row embedding)
     # feeds out-of-range token ids into gather and relies on backend
     # OOB behaviour, which is how we originally hit a :nan score.
-    @tag :distilbert_full
-    test "batched_run drives DistilBERT-QA through Nx.Serving" do
+    # `tag: :distilbert_full, lane_tags: false` keeps all three lanes gated
+    # behind `:distilbert_full` (not the lightweight `:native`), so they run
+    # under `--only distilbert_full` and never bloat `--only native`. The
+    # forward is driven through `Nx.Serving`'s `:defn_options`, which is
+    # where the compiler lanes plug in.
+    mode_test "batched_run drives DistilBERT-QA through Nx.Serving",
+      tag: :distilbert_full,
+      lane_tags: false do
       {:ok, model_info} =
         Bumblebee.load_model({:hf, "distilbert-base-uncased-distilled-squad"})
 
       {:ok, tokenizer} =
         Bumblebee.load_tokenizer({:hf, "distilbert-base-uncased-distilled-squad"})
 
-      serving = Bumblebee.Text.question_answering(model_info, tokenizer)
+      serving =
+        Bumblebee.Text.question_answering(model_info, tokenizer, defn_options: predict_opts)
 
       start_supervised!({Nx.Serving, serving: serving, name: __MODULE__.Serving})
 

test/emily/conformance/modernbert_test.exs

@@ -33,7 +33,7 @@ defmodule Emily.Conformance.ModernBertTest do
   @moduletag :conformance
   @moduletag capture_log: true
 
-  test "ModernBert :base forward on Emily.Backend" do
+  mode_test "ModernBert :base forward on Emily.Backend" do
     spec =
       Bumblebee.configure(ModernBert,
         architecture: :base,
@@ -47,7 +47,10 @@ defmodule Emily.Conformance.ModernBertTest do
       )
 
     model = ModernBert.model(spec)
-    {init_fn, predict_fn} = Axon.build(model)
+    # Init on the evaluator (params are random-init, mode-irrelevant);
+    # gate only the forward pass under `predict_opts`.
+    {init_fn, _} = Axon.build(model)
+    {_, predict_fn} = Axon.build(model, predict_opts)
 
     input_template = %{
       "input_ids" => Nx.template({1, 8}, :s64),

test/emily/conformance/nomic_embeddings_test.exs

@@ -27,7 +27,7 @@ defmodule Emily.Conformance.NomicEmbeddingsTest do
   @moduletag :conformance
   @moduletag capture_log: true
 
-  test "NomicBert :base forward runs end-to-end on Emily.Backend" do
+  mode_test "NomicBert :base forward runs end-to-end on Emily.Backend" do
     spec =
       Bumblebee.configure(NomicBert,
         architecture: :base,
@@ -41,7 +41,10 @@ defmodule Emily.Conformance.NomicEmbeddingsTest do
       )
 
     model = NomicBert.model(spec)
-    {init_fn, predict_fn} = Axon.build(model)
+    # Init on the evaluator (params are random-init, mode-irrelevant);
+    # gate only the forward pass under `predict_opts`.
+    {init_fn, _} = Axon.build(model)
+    {_, predict_fn} = Axon.build(model, predict_opts)
 
     input_template = %{
       "input_ids" => Nx.template({1, 8}, :s64),

test/emily/conformance/smollm3_test.exs

@@ -30,7 +30,7 @@ defmodule Emily.Conformance.SmolLm3Test do
   @moduletag :conformance
   @moduletag capture_log: true
 
-  test "SmolLm3 :for_causal_language_modeling forward on Emily.Backend" do
+  mode_test "SmolLm3 :for_causal_language_modeling forward on Emily.Backend" do
     spec =
       Bumblebee.configure(SmolLm3,
         architecture: :for_causal_language_modeling,
@@ -44,7 +44,10 @@ defmodule Emily.Conformance.SmolLm3Test do
       )
 
     model = SmolLm3.model(spec)
-    {init_fn, predict_fn} = Axon.build(model)
+    # Init on the evaluator (params are random-init, mode-irrelevant);
+    # gate only the forward pass under `predict_opts`.
+    {init_fn, _} = Axon.build(model)
+    {_, predict_fn} = Axon.build(model, predict_opts)
 
     input_template = %{"input_ids" => Nx.template({1, 8}, :s64)}
     params = init_fn.(input_template, Axon.ModelState.empty())

test/emily/conformance/vit_full_test.exs

@@ -30,7 +30,8 @@ defmodule Emily.Conformance.VitFullTest do
   @moduletag capture_log: true
   @moduletag timeout: 600_000
 
-  test "google/vit-base-patch16-224 forward pass matches pinned logits slice" do
+  mode_test "google/vit-base-patch16-224 forward pass matches pinned logits slice",
+    lane_tags: false do
     {:ok, %{model: model, params: params, spec: spec}} =
       Bumblebee.load_model({:hf, "google/vit-base-patch16-224"})
 
@@ -42,7 +43,7 @@ defmodule Emily.Conformance.VitFullTest do
       "pixel_values" => Nx.broadcast(Nx.tensor(0.5, type: :f32), {1, 224, 224, 3})
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.logits) == {1, 1000}
 

test/emily/conformance/vit_test.exs

@@ -31,7 +31,7 @@ defmodule Emily.Conformance.VitTest do
   @moduletag capture_log: true
   @moduletag timeout: 120_000
 
-  test ":base" do
+  mode_test ":base" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model({:hf, "hf-internal-testing/tiny-random-ViTModel"})
 
@@ -41,7 +41,7 @@ defmodule Emily.Conformance.VitTest do
       "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3})
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.hidden_state) == {1, 226, 32}
     assert Nx.shape(outputs.pooled_state) == {1, 32}
@@ -59,7 +59,7 @@ defmodule Emily.Conformance.VitTest do
     )
   end
 
-  test ":for_image_classification" do
+  mode_test ":for_image_classification" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(
                {:hf, "hf-internal-testing/tiny-random-ViTForImageClassification"}
@@ -71,7 +71,7 @@ defmodule Emily.Conformance.VitTest do
       "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3})
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.logits) == {1, 2}
 
@@ -81,7 +81,7 @@ defmodule Emily.Conformance.VitTest do
     )
   end
 
-  test ":for_masked_image_modeling" do
+  mode_test ":for_masked_image_modeling" do
     assert {:ok, %{model: model, params: params, spec: spec}} =
              Bumblebee.load_model(
                {:hf, "hf-internal-testing/tiny-random-ViTForMaskedImageModeling"}
@@ -93,7 +93,7 @@ defmodule Emily.Conformance.VitTest do
       "pixel_values" => Nx.broadcast(0.5, {1, 30, 30, 3})
     }
 
-    outputs = Axon.predict(model, params, inputs)
+    outputs = Axon.predict(model, params, inputs, predict_opts)
 
     assert Nx.shape(outputs.pixel_values) == {1, 30, 30, 3}