From 63b19d74a260cc069b76d927589ad683466b7902 Mon Sep 17 00:00:00 2001
From: Paulo Valente <16843419+polvalente@users.noreply.github.com>
Date: Tue, 24 Sep 2024 18:04:43 -0300
Subject: [PATCH 1/2] fix: test on stable ordering instead of ref-dependent
 order

---
 nx/test/nx/defn/tree_test.exs | 20 +++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/nx/test/nx/defn/tree_test.exs b/nx/test/nx/defn/tree_test.exs
index 6586119039d..fee1a77917f 100644
--- a/nx/test/nx/defn/tree_test.exs
+++ b/nx/test/nx/defn/tree_test.exs
@@ -1,5 +1,5 @@
 defmodule Nx.Defn.TreeTest do
-  use ExUnit.Case, async: true
+  use ExUnit.Case, async: false
 
   alias Nx.Defn.{Expr, Tree}
   doctest Nx.Defn.Tree
@@ -40,7 +40,9 @@ defmodule Nx.Defn.TreeTest do
 
     test "ignores constants" do
       a = Expr.parameter(:root, {:u, 64}, {}, 0)
-      assert [{_, :parameter}, {_, :add}] = plus_constant(a) |> Tree.scope_ids() |> Enum.sort()
+
+      assert [{_, :add}, {_, :parameter}] =
+               plus_constant(a) |> Tree.scope_ids() |> Enum.sort_by(&elem(&1, 1))
     end
 
     defn inside_cond(bool, a, b) do
@@ -54,8 +56,8 @@ defmodule Nx.Defn.TreeTest do
     test "ignores expressions inside cond" do
       {bool, cond} = Nx.Defn.jit(&{&1, inside_cond(&1, &2, &3)}).(0, 1, 2)
 
-      assert cond |> Tree.scope_ids() |> Enum.sort() ==
-               [{bool.data.id, :parameter}, {cond.data.id, :cond}]
+      assert cond |> Tree.scope_ids() |> Enum.sort_by(&elem(&1, 1)) ==
+               [{cond.data.id, :cond}, {bool.data.id, :parameter}]
     end
 
     defn inside_both_cond(bool, a, b) do
@@ -84,14 +86,14 @@ defmodule Nx.Defn.TreeTest do
       b = Expr.parameter(:root, {:u, 64}, {}, 2)
 
       assert [
-               {_, :parameter},
-               {_, :parameter},
-               {_, :parameter},
                {_, :add},
                {_, :cond},
                {_, :cond},
-               {_, :multiply}
-             ] = inside_both_cond(bool, a, b) |> Tree.scope_ids() |> Enum.sort()
+               {_, :multiply},
+               {_, :parameter},
+               {_, :parameter},
+               {_, :parameter}
+             ] = inside_both_cond(bool, a, b) |> Tree.scope_ids() |> Enum.sort_by(&elem(&1, 1))
     end
   end
 

From 7b531bbbd5fd7583676aa23ca8ff121470d12812 Mon Sep 17 00:00:00 2001
From: Paulo Valente <16843419+polvalente@users.noreply.github.com>
Date: Tue, 24 Sep 2024 18:05:13 -0300
Subject: [PATCH 2/2] Update nx/test/nx/defn/tree_test.exs

---
 nx/test/nx/defn/tree_test.exs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nx/test/nx/defn/tree_test.exs b/nx/test/nx/defn/tree_test.exs
index fee1a77917f..5ccd4fb3000 100644
--- a/nx/test/nx/defn/tree_test.exs
+++ b/nx/test/nx/defn/tree_test.exs
@@ -1,5 +1,5 @@
 defmodule Nx.Defn.TreeTest do
-  use ExUnit.Case, async: false
+  use ExUnit.Case, async: true
 
   alias Nx.Defn.{Expr, Tree}
   doctest Nx.Defn.Tree