Lower Nx.top_k (Nx.Block.TopK) in the native Expr compiler

RELEASE.md

@@ -75,6 +75,14 @@
   / `mx::scatter_add` (accumulate) bit-for-bit. Layouts MLX can't scatter
   still route through the evaluator under `native_fallback: :eval`.
 
+- **`Nx.top_k` lowers natively.** The `Nx.Block.TopK` block — a multi-output
+  `{values, indices}` — now compiles instead of raising. `Emily.Backend` has
+  no `top_k` override (`mx::topk` yields values only, not the indices Nx's
+  contract requires), so the evaluator computes it via the block's default
+  expansion (`argsort(desc)` → `take_along_axis` → `slice` the top k); the
+  compiler lowers that same expansion — every op in it already lowers — and
+  projects the two leaves via `:elem`, bit-identical to the evaluator.
+
 - **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

lib/emily/ir.ex

@@ -1145,6 +1145,36 @@ defmodule Emily.IR do
     emit_coerced(state, :irfftn, [rt], [[length], [axis]], out.type)
   end
 
+  # Nx.top_k (Nx.Block.TopK) — a multi-output block `{values, indices}`.
+  # Emily.Backend has no `top_k` override (mx::topk yields values only, not the
+  # indices Nx's contract requires), so the evaluator computes it via the
+  # block's *default expansion*: argsort(desc) -> take_along_axis -> slice the
+  # top k. Every op in that expansion already lowers, so lower the expansion
+  # itself and hand the two leaf refs back as a projectable tuple for `:elem`
+  # (the `{:multi_refs, …}` machinery the tuple-`cond` path added). Because it
+  # IS the evaluator's expansion, the result is bit-identical.
+  defp lower_block(%Nx.Block.TopK{}, in_args, expr, _t, state) when is_tuple(expr) do
+    # `Nx.Defn.Expr.expr_block` builds the expansion against FRESH per-position
+    # `:parameter` nodes, not the real `in_args` — so lowering `expr` directly
+    # would resolve a block parameter to the OUTER function's `{:input, pos}`
+    # (a different tensor entirely). Bind them first: lower the real args, then
+    # seed the cache so each block parameter resolves to its arg's ref (the
+    # same binding `while` does via `param_seed`). Then lower the expansion's
+    # leaves — bit-identical to the evaluator, which runs the same expansion.
+    args = Enum.take_while(in_args, &(not is_list(&1)))
+    {arg_refs, state} = Enum.map_reduce(args, state, &lower_node/2)
+    leaves = Tuple.to_list(expr)
+
+    seed =
+      leaves
+      |> Enum.reduce(%{}, &collect_block_params/2)
+      |> Map.new(fn {id, pos} -> {id, Enum.at(arg_refs, pos)} end)
+
+    state = %{state | cache: Map.merge(state.cache, seed)}
+    {leaf_refs, state} = Enum.map_reduce(leaves, state, &lower_node/2)
+    {{:multi_refs, leaf_refs}, state}
+  end
+
   # Any other block struct raises. Lowering the block's composed
   # expansion would silently diverge from the Evaluator whenever
   # Emily.Backend.block/4 dispatches that struct through a fused / native
@@ -1208,6 +1238,25 @@ defmodule Emily.IR do
 
   defp tensors?(list), do: Enum.all?(list, &match?(%T{}, &1))
 
+  # Collect `{parameter_id => position}` for every block-local `:parameter`
+  # reachable from `node`. Used to bind a block expansion's fresh parameters
+  # to the real in_args (see the Nx.Block.TopK lowering). Walks only the small
+  # expansion graph — parameters are leaves, so it never descends into the
+  # in_args' own (possibly large) graph.
+  defp collect_block_params(%T{data: %Nx.Defn.Expr{op: :parameter, id: id, args: [pos]}}, acc),
+    do: Map.put_new(acc, id, pos)
+
+  defp collect_block_params(%T{data: %Nx.Defn.Expr{args: args}}, acc),
+    do: Enum.reduce(args, acc, &collect_block_params/2)
+
+  defp collect_block_params(list, acc) when is_list(list),
+    do: Enum.reduce(list, acc, &collect_block_params/2)
+
+  defp collect_block_params(tuple, acc) when is_tuple(tuple),
+    do: Enum.reduce(Tuple.to_list(tuple), acc, &collect_block_params/2)
+
+  defp collect_block_params(_other, acc), do: acc
+
   defp float_like?({kind, _}) when kind in [:f, :bf, :c], do: true
   defp float_like?(_), do: false
 

test/emily/compiler_equivalence_test.exs

@@ -502,6 +502,47 @@ defmodule Emily.CompilerEquivalenceTest do
     end
   end
 
+  describe "top_k (multi-output block)" do
+    test "values and indices both match the evaluator" do
+      x = et([[3.0, 1.0, 4.0, 1.5, 5.0], [9.0, 2.0, 6.0, 5.0, 3.0]])
+
+      {nv, ni} = run(fn t -> Nx.top_k(t, k: 3) end, [x], @native)
+      {ev, ei} = run(fn t -> Nx.top_k(t, k: 3) end, [x], @eval)
+
+      assert %Emily.Backend{} = nv.data
+      assert Nx.to_binary(nv) == Nx.to_binary(ev)
+      assert Nx.to_binary(ni) == Nx.to_binary(ei)
+      assert ni.type == {:s, 32}
+    end
+
+    test "top_k feeding a downstream op lowers (both leaves projected)" do
+      x = et([[3.0, 1.0, 4.0, 1.5, 5.0]])
+
+      assert_equiv(
+        fn t ->
+          {v, i} = Nx.top_k(t, k: 2)
+          Nx.add(v, Nx.as_type(i, :f32))
+        end,
+        [x]
+      )
+    end
+
+    test "top_k binds its block param to the real arg, not outer {:input, 0} (regression)" do
+      # Mirrors the ModernBERT failure: outer input 0 is lower-rank than the
+      # top_k input. The `Nx.Block.TopK` expansion is built against fresh
+      # parameters, so an unbound parameter falls through to `{:input, 0}` —
+      # here the 1-D tensor — and argsort hits "axis 1 on a 1-D array".
+      one_d = et([0.0, 0.0])
+      scores = et([[3.0, 1.0, 4.0], [1.0, 5.0, 9.0]])
+
+      {nv, ni} = run(fn _i, s -> Nx.top_k(s, k: 2) end, [one_d, scores], @native)
+      {ev, ei} = run(fn _i, s -> Nx.top_k(s, k: 2) end, [one_d, scores], @eval)
+
+      assert Nx.to_binary(nv) == Nx.to_binary(ev)
+      assert Nx.to_binary(ni) == Nx.to_binary(ei)
+    end
+  end
+
   describe "fft family (signal transforms)" do
     test "1-D fft / ifft on the trailing axis match the evaluator" do
       x = et([1.0, 2.0, 3.0, 4.0])