feat: clear errors for unsupported control flow (reduce fn / while)

- reduce / window_reduce with an arbitrary BEAM reducer raise a clear
  error (no silent fallback); the fixed-identity aggregates (sum/product/
  reduce_max/reduce_min) lower natively as their own ops.
- while and its :elem tuple projection raise a clear "deferred" error —
  the single-NIF replay has no loop construct, so a data-dependent while
  needs static-trip unrolling or a worker-side synced loop (a focused
  follow-up). defn while is not used by the core transformer forwards;
  decode/generation loops run in Elixir.

Tests assert both raise clearly under the native compiler.

Author: ausimian

lib/emily/ir.ex

@@ -195,6 +195,7 @@ defmodule Emily.IR do
 
   alias Emily.Fast.Block, as: FB
   alias Emily.Quantization.Block, as: QB
+  alias Nx.Defn.Tree
   alias Nx.Tensor, as: T
 
   # Nx Expr op -> IR opcode. Arithmetic/bitwise cast both operands to the
@@ -547,7 +548,7 @@ defmodule Emily.IR do
   # hook the token is a no-op, so pass through to the inner expr. Hooks
   # would need a callback into Elixir mid-graph (program-split) — deferred.
   defp lower_op(%T{data: %Nx.Defn.Expr{op: :attach_token, args: [token, expr]}}, state) do
-    if Nx.Defn.Tree.has_hooks?(token, %{}) do
+    if Tree.has_hooks?(token, %{}) do
       raise ArgumentError,
             "Emily Expr compiler does not support hooks under native compilation " <>
               "(they require a mid-graph callback into Elixir)."
@@ -556,6 +557,27 @@ defmodule Emily.IR do
     lower_node(expr, state)
   end
 
+  # reduce / window_reduce with a user-supplied BEAM reducer cannot be
+  # compiled — the reducer would have to run on the host mid-graph. The
+  # fixed-identity aggregates (sum/product/max/min) are separate ops and
+  # already lower natively; only an arbitrary reducer reaches here.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: op}}, _state) when op in [:reduce, :window_reduce] do
+    raise ArgumentError,
+          "Emily Expr compiler cannot lower #{inspect(op)} with an arbitrary " <>
+            "reducer function (it would require a host callback mid-graph; no " <>
+            "fallback). Use the native aggregates (sum/product/reduce_max/" <>
+            "reduce_min) where possible."
+  end
+
+  # while / its tuple projection are deferred to a follow-up: the
+  # single-NIF replay has no loop construct, so a data-dependent while
+  # needs either static-trip unrolling or a worker-side synced loop.
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: op}}, _state) when op in [:while, :elem] do
+    raise ArgumentError,
+          "Emily Expr compiler does not yet lower #{inspect(op)} (defn `while` " <>
+            "loops; deferred — the decode/generation loops run in Elixir today)."
+  end
+
   defp lower_op(%T{data: %Nx.Defn.Expr{op: op}}, _state) do
     raise ArgumentError,
           "Emily Expr compiler does not yet lower op #{inspect(op)} " <>

test/emily/compiler_control_flow_test.exs

@@ -33,6 +33,15 @@ defmodule Emily.CompilerControlFlowTest do
     if Nx.greater(Nx.reduce_max(y), 5), do: Nx.divide(y, 2), else: y
   end
 
+  defn while_fn(x) do
+    {_i, acc} =
+      while {i = 0, acc = x}, i < 5 do
+        {i + 1, Nx.multiply(acc, 2)}
+      end
+
+    acc
+  end
+
   defp equiv(fun, x) do
     native = Nx.Defn.jit(fun, @native).(x)
     eval = Nx.Defn.jit(fun, @eval).(x)
@@ -60,4 +69,20 @@ defmodule Emily.CompilerControlFlowTest do
       end
     end
   end
+
+  describe "unsupported control flow raises (no silent fallback)" do
+    test "arbitrary reduce/2 fn raises a clear error" do
+      f = fn x -> Nx.reduce(x, 0.0, fn a, b -> Nx.add(a, b) end) end
+
+      assert_raise ArgumentError, ~r/arbitrary reducer/, fn ->
+        Nx.Defn.jit(f, @native).(Nx.tensor([1.0, 2.0, 3.0], backend: Emily.Backend))
+      end
+    end
+
+    test "defn while raises (deferred)" do
+      assert_raise ArgumentError, ~r/while/, fn ->
+        Nx.Defn.jit(&while_fn/1, @native).(Nx.tensor([1.0, 2.0], backend: Emily.Backend))
+      end
+    end
+  end
 end