feat: lower indexed_put / indexed_add in the native Expr compiler

Nx.indexed_put and Nx.indexed_add (scatter) now lower under
`compiler: Emily.Compiler, native: true` for MLX-scatter-compatible index
layouts, instead of forcing a graceful fallback to the evaluator. The MLX
ops and eager NIFs already existed (Native.{scatter,scatter_add} ->
mx::{scatter,scatter_add}); only the compiler path was missing.

- opcodes.hpp: add Scatter/ScatterAdd (92-93), bump kOpcodeCount to 94,
  dispatch to mx::scatter (overwrite) / mx::scatter_add (accumulate).
  operands [target, updates, idx0, ...]; iattrs [[axes...]].
- ir.ex: add the opcodes; lower :indexed_put/:indexed_add (one shared
  clause), mirroring Emily.Backend's apply_scatter — reuse the existing
  gather index-split + scatter_gather_compatible? helpers, port
  updates_shape_for_scatter. Incompatible index layouts raise (no
  fallback), matching the native gather; the evaluator handles them under
  native_fallback: :eval.
- compiler_equivalence_test.exs: native-vs-evaluator bit-identical cases
  for indexed_put/indexed_add (2-D grid, duplicate-index accumulation,
  partial-axis whole-row writes).

Another op surfaced while running the native compiler on the Whisper
livebook.

Author: ausimian

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

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
@@ -1285,4 +1323,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_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.