fix: while with vectorized cond inside (#1474)

Author: polvalente

exla/test/exla/defn/expr_test.exs

@@ -1465,6 +1465,38 @@ defmodule EXLA.Defn.ExprTest do
       assert_equal(while_inside_cond(1, 1), 1)
       assert_equal(while_inside_cond(1, 0), 2)
     end
+
+    defn cond_inside_while_vectorized(t, size) do
+      down = Nx.u8(0)
+      up = Nx.u8(1)
+      mode = down
+      i = Nx.s64(0)
+
+      [t, node, i, size, mode] =
+        Nx.broadcast_vectors([t, 0, i, size, mode])
+
+      {t, _} =
+        while {t, {node, i, _mode = mode, size, up}},
+              node != -1 and i >= 0 do
+          mode =
+            cond do
+              node >= size -> up
+              true -> -1
+            end
+
+          {t, {node, i - 1, mode, size, up}}
+        end
+
+      t
+    end
+
+    test "cond inside vectorized while" do
+      assert_raise CompileError,
+                   ~r/the do-block in while must return tensors with the same shape, type, and names as the initial arguments./,
+                   fn ->
+                     cond_inside_while_vectorized(Nx.vectorize(Nx.tensor([1, 2, 3]), :a), 3)
+                   end
+    end
   end
 
   describe "map" do

nx/lib/nx/defn/expr.ex

@@ -191,7 +191,7 @@ defmodule Nx.Defn.Expr do
     last
   end
 
-  defp non_vectorized_cond(clauses, last = out, context) do
+  defp non_vectorized_cond(clauses, out_template = last, context) do
     {preds, exprs} = Enum.unzip(clauses)
 
     {broadcasted_clauses, vectorized_axes} =
@@ -211,10 +211,10 @@ defmodule Nx.Defn.Expr do
 
     out =
       if vectorized_axes == [] do
-        out
+        out_template
       else
         {result, []} =
-          Composite.traverse(out, vectorized_axes, fn
+          Composite.traverse(out_template, vectorized_axes, fn
             %T{} = expr, [axes | tail] ->
               {%{expr | vectorized_axes: axes}, tail}
 
@@ -538,6 +538,9 @@ defmodule Nx.Defn.Expr do
         {inner_arg, inner_context} = to_param_expr(inner_initial, :while)
         inner_condition = condition_body.(:condition, inner_arg) |> to_pred(line, file, :while)
         inner_body = condition_body.(:body, inner_arg) |> to_container_expr()
+
+        compatible_while!(file, line, inner_initial, inner_body)
+
         inner_while = while(inner_initial, inner_context, inner_arg, inner_condition, inner_body)
 
         vectorized_while__build_outer_while(