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

[ausimian]

Nx.top_k now compiles under native instead of raising "does not yet lower
the block Nx.Block.TopK". With this, all 8 example livebooks lower fully
native — modernbert_classification was the lone holdout (found by the
forced-native livebook probe).

What it is

Nx.Block.TopK is a multi-output block ({values, indices}), and
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 there already lowers, so the compiler lowers that same
expansion and projects the two leaves via :elem (the {:multi_refs, …}
machinery the tuple-cond fix added) — bit-identical to the evaluator.

The subtle part — binding the block's parameters

Nx.Defn.Expr.expr_block builds a block's expansion against fresh
per-position :parameter nodes, not the real in_args. So lowering the
expansion directly makes a block parameter 0 fall through to the outer
function's {:input, 0} — a completely different tensor. That's exactly how
ModernBERT failed: its outer input 0 is the 1-D input_ids, so top_k's
argsort ran on a 1-D array → "invalid axis 1 for array with 1 dimensions"
at replay, even though logits/softmax were correct {1, num_labels}.

The fix binds the block parameters to the lowered in_args (the same thing
while does via param_seed) before lowering the expansion. Eval never tripped
on this because it evaluates the bound expansion directly on real tensors.

Changes

• ir.ex — lower_block for Nx.Block.TopK: lower the in_args, seed the
  cache so each block :parameter resolves to its arg, lower the expansion
  leaves, return {:multi_refs, [vref, iref]}. collect_block_params/2 walks
  the (small) expansion to find the parameters.
• compiler_equivalence_test.exs — values + s32 indices match; a
  downstream op consuming both leaves; and a regression case where outer
  input 0 is 1-D (the ModernBERT shape), which an unbound parameter would crash.

Verification

• mix precommit green: 694 tests, 0 failures.
• The modernbert_classification livebook, forced to native: true, native_fallback: :raise, now runs end-to-end native (was the 1 failing of
  8). The other 7 were already native.