fix: lower tuple-returning cond (multi-output) natively

A `cond`/`if` whose branches return a tuple crashed the native compiler
with a FunctionClauseError: the :cond lowering passed the tuple branch
straight to lower_node/2, which only matches a single tensor. Worse, a
FunctionClauseError escapes the compiler's graceful-fallback rescue (which
only traps ArgumentError), so it hard-faulted instead of degrading to the
evaluator. Surfaced by a Whisper encoder forward under native: true
(transpose(elem(cond -> {f32[1,1500,6,64] x N}, i))).

Lower a tuple cond to one where-chain per leaf position — identical
wholesale-select semantics to the single-output cond (the predicate is a
whole-tensor scalar bool; every branch is still computed) — and return a
{:multi_refs, [...]} handle that :elem projects, shared across sibling
:elems via lower_node's memo (like while). A nested / non-tensor container
raises a clean ArgumentError, so it falls back gracefully rather than
crashing.

- ir.ex: split the :cond clause (single-tensor head unchanged; new
  is_tuple head for the multi-output case); extend :elem to project
  {:multi_refs, refs}; add the tensors?/1 guard.
- compiler_control_flow_test.exs: tuple if + multi-clause tuple cond with
  mixed-shape (vector + scalar) leaves, native-vs-evaluator bit-identical.

Author: ausimian

lib/emily/ir.ex

@@ -870,7 +870,7 @@ defmodule Emily.IR do
   # gather/index there clamps rather than faults, so the discarded value
   # never changes the result; a hard-faulting op on a not-taken path would
   # diverge from the Evaluator's lazy single-branch eval.
-  defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, last]}} = t, state) do
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, %T{} = last]}} = t, state) do
     {last_ref, state} = lower_node(last, state)
 
     {result, state} =
@@ -884,6 +884,52 @@ defmodule Emily.IR do
     coerce(result, t.type, state)
   end
 
+  # Multi-output cond: each branch returns a tuple of tensors. Lower to one
+  # `where`-chain per leaf position — identical wholesale-select semantics to
+  # the single-output case above (the predicate is a whole-tensor scalar bool;
+  # every branch is still computed). Returns a `{:multi_refs, [...]}` handle
+  # that `:elem` projects; sibling `:elem`s share it via lower_node's memo.
+  # A nested / non-tensor container raises (no fallback path for it yet).
+  defp lower_op(%T{data: %Nx.Defn.Expr{op: :cond, args: [clauses, last]}}, state)
+       when is_tuple(last) do
+    last_leaves = Tuple.to_list(last)
+
+    unless tensors?(last_leaves) and
+             Enum.all?(clauses, fn {_p, b} -> is_tuple(b) and tensors?(Tuple.to_list(b)) end) do
+      raise ArgumentError,
+            "Emily Expr compiler: cond over a nested / non-tensor container is not " <>
+              "lowered yet (only a flat tuple of tensors)."
+    end
+
+    # Lower each predicate (cast to pred) + its branch's leaf refs once, so the
+    # per-leaf where-chains share them.
+    {clauses, state} =
+      Enum.map_reduce(clauses, state, fn {pred, body}, st ->
+        {pred_ref, st} = lower_node(pred, st)
+        {pred_ref, st} = emit(st, :astype, [pred_ref], [[dtype_code({:pred, 1})]])
+        {body_refs, st} = Enum.map_reduce(Tuple.to_list(body), st, &lower_node/2)
+        {{pred_ref, body_refs}, st}
+      end)
+
+    rev = Enum.reverse(clauses)
+
+    {refs, state} =
+      last_leaves
+      |> Enum.with_index()
+      |> Enum.map_reduce(state, fn {leaf, j}, st ->
+        {last_ref, st} = lower_node(leaf, st)
+
+        {result, st} =
+          Enum.reduce(rev, {last_ref, st}, fn {pred_ref, body_refs}, {else_ref, st2} ->
+            emit(st2, :where, [pred_ref, Enum.at(body_refs, j), else_ref])
+          end)
+
+        coerce(result, leaf.type, st)
+      end)
+
+    {{:multi_refs, refs}, state}
+  end
+
   # attach_token: sequences a token (hooks) before `expr`. With no active
   # 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.
@@ -955,11 +1001,14 @@ defmodule Emily.IR do
       {{:multi, base, _arity}, state} ->
         {{:instr, base + i}, state}
 
+      {{:multi_refs, refs}, state} ->
+        {Enum.at(refs, i), state}
+
       {_handle, _state} ->
         raise ArgumentError,
               "Emily Expr compiler: :elem projects a tuple-producing op it can't " <>
-                "lower yet (only `while` produces projectable tuples today; other " <>
-                "multi-output ops are unsupported)."
+                "lower yet (only `while` and tuple `cond` produce projectable " <>
+                "tuples today; other multi-output ops are unsupported)."
     end
   end
 
@@ -1157,6 +1206,8 @@ defmodule Emily.IR do
   defp bool_int(true), do: 1
   defp bool_int(false), do: 0
 
+  defp tensors?(list), do: Enum.all?(list, &match?(%T{}, &1))
+
   defp float_like?({kind, _}) when kind in [:f, :bf, :c], do: true
   defp float_like?(_), do: false
 

test/emily/compiler_control_flow_test.exs

@@ -45,6 +45,34 @@ defmodule Emily.CompilerControlFlowTest do
     acc
   end
 
+  # Tuple-returning cond (multi-output): both branches produce a {vec, vec}
+  # tuple; `:elem` projects each leaf out of the lowered per-leaf where-chains.
+  defn tuple_if_fn(x) do
+    {a, b} =
+      if Nx.greater(Nx.sum(x), 0) do
+        {Nx.multiply(x, 2), Nx.add(x, 1)}
+      else
+        {Nx.negate(x), Nx.subtract(x, 1)}
+      end
+
+    Nx.add(a, b)
+  end
+
+  # Multi-clause tuple cond with leaves of DIFFERENT shapes (vector + scalar),
+  # so each leaf's where-chain is built and coerced independently.
+  defn tuple_cond3_fn(x) do
+    s = Nx.sum(x)
+
+    {vec, scalar} =
+      cond do
+        Nx.greater(s, 10) -> {Nx.multiply(x, 10), Nx.reduce_max(x)}
+        Nx.greater(s, 0) -> {Nx.multiply(x, 2), Nx.product(x)}
+        true -> {Nx.negate(x), s}
+      end
+
+    Nx.add(vec, scalar)
+  end
+
   defp equiv(fun, x) do
     native = Nx.Defn.jit(fun, @native).(x)
     eval = Nx.Defn.jit(fun, @eval).(x)
@@ -71,6 +99,18 @@ defmodule Emily.CompilerControlFlowTest do
         equiv(&nested_if_fn/1, Nx.tensor(data, backend: Emily.Backend))
       end
     end
+
+    test "tuple-returning if (multi-output cond) projects each leaf" do
+      for data <- [[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]] do
+        equiv(&tuple_if_fn/1, Nx.tensor(data, backend: Emily.Backend))
+      end
+    end
+
+    test "multi-clause tuple cond with mixed-shape leaves matches the evaluator" do
+      for data <- [[5.0, 4.0, 3.0], [1.0, 1.0, 1.0], [-2.0, -2.0, -2.0]] do
+        equiv(&tuple_cond3_fn/1, Nx.tensor(data, backend: Emily.Backend))
+      end
+    end
   end
 
   describe "defn while" do