Skip to content

Commit

Permalink
fix: test on stable ordering instead of ref-dependent order
Browse files Browse the repository at this point in the history
  • Loading branch information
polvalente committed Sep 24, 2024
1 parent 93e4383 commit 63b19d7
Showing 1 changed file with 11 additions and 9 deletions.
20 changes: 11 additions & 9 deletions nx/test/nx/defn/tree_test.exs
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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

Expand Down

0 comments on commit 63b19d7

Please sign in to comment.