Rename notebooks/ → livebooks/

README.md

@@ -191,16 +191,16 @@ Nx.Defn.global_default_options(compiler: Emily.Compiler)
 Bumblebee inference works with no further configuration once the
 backend is installed — see the conformance suites under
 `test/emily/conformance/` for worked DistilBERT, Qwen3, ViT, and
-Whisper pipelines, and the Notebooks section of the HexDocs nav for
+Whisper pipelines, and the Livebooks section of the HexDocs nav for
 runnable Livebooks.
 
 The low-level tensor API (`Emily.from_binary/3`, `to_binary/1`,
 `shape/1`, `dtype/1`, `eval/1`) remains available for diagnostics
 and direct MLX round-trips, but most users should go through Nx.
 
-## Notebooks
+## Livebooks
 
-End-to-end Livebooks under `notebooks/`. Each one declares its own
+End-to-end Livebooks under `livebooks/`. Each one declares its own
 `Mix.install/2` block and pins `Emily.Backend` as the default Nx
 backend, so they're self-contained — open in Livebook and run.
 
@@ -339,8 +339,8 @@ caller-facing API; the only difference is whether the calling
 process wraps the call in `Emily.Stream.with_stream/2` and whether
 you run one serving or many.
 
-See `Emily.Stream` for the API and the `qwen3_quantized` notebook
-under Notebooks for a worked multi-stream example.
+See `Emily.Stream` for the API and the `qwen3_quantized` livebook
+under Livebooks for a worked multi-stream example.
 
 ## Observability
 
@@ -438,7 +438,7 @@ be introduced in the layer where its test fails.
 ## Documentation
 
   * [HexDocs](https://hexdocs.pm/emily) — per-module API docs and
-    runnable notebooks.
+    runnable livebooks.
   * [`ARCHITECTURE.md`](ARCHITECTURE.md) — current shape of the
     library: layer boundaries, design decisions, concurrency and
     memory model, observability surface.

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`

ROADMAP.md

@@ -74,5 +74,5 @@ Tracking checklist:
   * `CHANGELOG.md` accumulated across releases (it is, since 0.3.0).
   * `MAINTAINING.md` reflects the precompiled-NIF release flow (it
     does, since 0.3.0).
-  * Worked Bumblebee + quantized-Qwen3 examples in `notebooks/`
-    (present and grouped in the HexDocs Notebooks section).
+  * Worked Bumblebee + quantized-Qwen3 examples in `livebooks/`
+    (present and grouped in the HexDocs Livebooks section).

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

notebooks/mnist_training.livemd → livebooks/mnist_training.livemd

@@ -17,7 +17,7 @@ Mix.install(
 
 ## Overview
 
-The other notebooks in this repo run pretrained models forward. This
+The other livebooks in this repo run pretrained models forward. This
 one trains from scratch — a small Axon MLP on MNIST, lowered through
 `Emily.Compiler` so every step of the backward pass dispatches to
 MLX. It's the end-to-end exercise of the `Nx.Defn.grad` chain that

mix.exs

@@ -230,14 +230,14 @@ defmodule Emily.MixProject do
         "ARCHITECTURE.md",
         "ROADMAP.md",
         "CHANGELOG.md",
-        "notebooks/distilbert_qa.livemd",
-        "notebooks/qwen3_quantized.livemd",
-        "notebooks/nomic_embeddings.livemd",
-        "notebooks/smollm3_chat.livemd",
-        "notebooks/modernbert_classification.livemd",
-        "notebooks/mnist_training.livemd",
-        "notebooks/whisper_transcription.livemd",
-        "notebooks/fast_kernels.livemd"
+        "livebooks/distilbert_qa.livemd",
+        "livebooks/qwen3_quantized.livemd",
+        "livebooks/nomic_embeddings.livemd",
+        "livebooks/smollm3_chat.livemd",
+        "livebooks/modernbert_classification.livemd",
+        "livebooks/mnist_training.livemd",
+        "livebooks/whisper_transcription.livemd",
+        "livebooks/fast_kernels.livemd"
       ],
       groups_for_extras: [
         README: ~r{README.md},
@@ -246,7 +246,7 @@ defmodule Emily.MixProject do
           "ROADMAP.md",
           "CHANGELOG.md"
         ],
-        Notebooks: ~r{^notebooks/}
+        Livebooks: ~r{^livebooks/}
       ],
       groups_for_modules: [
         Core: [Emily, Emily.Backend, Emily.Compiler],

scripts/test-livebooks.sh

@@ -1,7 +1,7 @@
 #!/usr/bin/env bash
 # Smoke-test the example livebooks against the LOCAL emily checkout.
 #
-# For each notebook in notebooks/, this extracts its Elixir cells,
+# For each notebook in livebooks/, this extracts its Elixir cells,
 # repoints the `{:emily, "~> x"}` Mix.install dependency at this repo (as
 # a `path:` dep, so the notebook exercises the working tree — including a
 # from-source NIF build), and runs the result headlessly with `elixir`.
@@ -28,7 +28,7 @@
 set -uo pipefail
 
 REPO="$(cd "$(dirname "$0")/.." && pwd)"
-NB_DIR="$REPO/notebooks"
+NB_DIR="$REPO/livebooks"
 TIMEOUT="${LIVEBOOK_TIMEOUT:-1200}"
 SKIP=" ${LIVEBOOK_SKIP:-} "
 

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