Add runtime_call regression tests for while, cond, and inline fns

- runtime_call inside while loop (increment to 10)
- while loop with tuple state (multiply by 2 three times)
- cond branches (positive → double, negative → negate)
- multiple runtime_calls in one while body (add 1 then double)
- type-changing callback (s32 → f32)
- nested while loops with runtime_call
- while with separate accumulator
- tuple input inside while (skipped: shape mismatch on EXLA,
  wrong result on evaluator)

All 17 tests pass (1 skipped). Tested on EXLA host.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: blasphemetheus

exla/test/exla/defn/runtime_call_test.exs

@@ -209,14 +209,135 @@ defmodule EXLA.Defn.RuntimeCallTest do
     assert_receive {:container_fun, ^ref}
   end
 
+  def add_one_callback(t, _opts), do: Nx.add(t, 1)
+  def double_callback(t, _opts), do: Nx.multiply(t, 2)
+  def negate_callback(t, _opts), do: Nx.negate(t)
+
   defn runtime_call_in_while(x) do
     while x, Nx.less(x, 10) do
-      Nx.runtime_call(x, x, fn t -> Nx.add(t, 1) end)
+      Nx.runtime_call(x, x, &add_one_callback/2)
     end
   end
 
   test "runtime_call inside while loop" do
     result = runtime_call_in_while(Nx.tensor(0))
     assert_equal(result, Nx.tensor(10))
   end
+
+  defn runtime_call_in_while_with_tuple(x) do
+    {result, _count} =
+      while {x, count = Nx.tensor(0)}, Nx.less(count, 3) do
+        doubled = Nx.runtime_call(x, x, &double_callback/2)
+        {doubled, count + 1}
+      end
+
+    result
+  end
+
+  test "runtime_call inside while loop with tuple state" do
+    result = runtime_call_in_while_with_tuple(Nx.tensor(1.0))
+    # 1.0 * 2 * 2 * 2 = 8.0
+    assert_equal(result, Nx.tensor(8.0))
+  end
+
+  defn runtime_call_in_cond(x) do
+    if Nx.greater(x, 0) do
+      Nx.runtime_call(x, x, &double_callback/2)
+    else
+      Nx.runtime_call(x, x, &negate_callback/2)
+    end
+  end
+
+  test "runtime_call inside cond branches" do
+    assert_equal(runtime_call_in_cond(Nx.tensor(5.0)), Nx.tensor(10.0))
+    assert_equal(runtime_call_in_cond(Nx.tensor(-3.0)), Nx.tensor(3.0))
+  end
+
+  defn multiple_runtime_calls_in_while(x) do
+    while x, Nx.less(x, 100) do
+      step1 = Nx.runtime_call(x, x, &add_one_callback/2)
+      Nx.runtime_call(step1, step1, &double_callback/2)
+    end
+  end
+
+  test "multiple runtime_calls in one while body" do
+    # (0+1)*2=2, (2+1)*2=6, (6+1)*2=14, (14+1)*2=30, (30+1)*2=62, (62+1)*2=126
+    result = multiple_runtime_calls_in_while(Nx.tensor(0.0))
+    assert_equal(result, Nx.tensor(126.0))
+  end
+
+  def cast_to_float_callback(t, _opts), do: Nx.as_type(t, :f32)
+
+  defn runtime_call_type_change(x) do
+    out = %{x | type: {:f, 32}}
+    Nx.runtime_call(out, x, &cast_to_float_callback/2)
+  end
+
+  test "runtime_call where callback changes type" do
+    result = runtime_call_type_change(Nx.tensor([1, 2, 3], type: :s32))
+    assert Nx.type(result) == {:f, 32}
+    assert_equal(result, Nx.tensor([1.0, 2.0, 3.0]))
+  end
+
+  def add_ten_callback(t, _opts), do: Nx.add(t, 10)
+
+  defn nested_while_with_runtime_call(x) do
+    {result, _} =
+      while {x, outer = Nx.tensor(0)}, Nx.less(outer, 2) do
+        {inner_result, _} =
+          while {x, inner = Nx.tensor(0)}, Nx.less(inner, 3) do
+            {Nx.runtime_call(x, x, &add_one_callback/2), inner + 1}
+          end
+
+        {inner_result, outer + 1}
+      end
+
+    result
+  end
+
+  test "runtime_call inside nested while loops" do
+    result = nested_while_with_runtime_call(Nx.tensor(0.0))
+    # Inner while runs 3 times each outer iteration, outer runs 2 times
+    # 0 → +3 = 3 → +3 = 6
+    assert_equal(result, Nx.tensor(6.0))
+  end
+
+
+  def sum_tuple_callback({a, b}, _opts), do: Nx.add(a, b)
+
+  defn runtime_call_tuple_in_while(x, y) do
+    {result, _} =
+      while {x, count = Nx.tensor(0)}, Nx.less(count, 3) do
+        summed = Nx.runtime_call(x, {x, y}, &sum_tuple_callback/2)
+        {summed, count + 1}
+      end
+
+    result
+  end
+
+  @tag :skip
+  test "runtime_call with tuple input inside while" do
+    # Crashes on EXLA with shape_mismatch in args_spec.
+    # Returns wrong value (4.0 instead of 31.0) on Evaluator.
+    # Tuple inputs to runtime_call inside while with captured
+    # variables don't work correctly.
+    result = runtime_call_tuple_in_while(Nx.tensor(1.0), Nx.tensor(10.0))
+    assert_equal(result, Nx.tensor(31.0))
+  end
+
+  defn runtime_call_in_while_accumulating(x) do
+    {_x, acc} =
+      while {x, acc = Nx.tensor(0.0)}, Nx.less(acc, 100) do
+        val = Nx.runtime_call(x, x, &double_callback/2)
+        {val, acc + val}
+      end
+
+    acc
+  end
+
+  test "runtime_call in while with separate accumulator" do
+    # x=5, iter 1: val=10, acc=10; iter 2: val=20, acc=30; iter 3: val=40, acc=70; iter 4: val=80, acc=150
+    result = runtime_call_in_while_accumulating(Nx.tensor(5.0))
+    assert_equal(result, Nx.tensor(150.0))
+  end
 end