refactor(test): promote window_scatter_min/window_scatter_max tests to nx_test.exs (#85)

* window_scatter_min not working

* added window_scatter_min with vectorized input

* Updated window_scatter_min and max to feature in nx_test.exs; removed comments regarding the tie breaker in emlx.ex

* fix: simplify pad invocations

---------

Co-authored-by: Paulo Valente <16843419+polvalente@users.noreply.github.com>

Author: Chapaman

lib/emlx.ex

@@ -177,7 +177,7 @@ defmodule EMLX do
   deftensor tensordot(tensorA, tensorB, axesA, axesB)
   deftensor einsum(tensorA, tensorB, spec_string)
   deftensor transpose(tensor, axes)
-  deftensor pad(tensor, axes, low_pad_size, high_pad_size, pad_value)
+  deftensor pad(tensor, axes, low_pad_size, high_pad_size, tensor_pad_value)
   deftensor sort(tensor, axis)
   deftensor argsort(tensor, axis)
   deftensor tri_inv(tensor, upper)
@@ -449,9 +449,6 @@ defmodule EMLX do
   @impl Nx.Defn.Compiler
   defdelegate __partitions_options__(opts), to: Nx.Defn.Evaluator
 
-  @impl Nx.Defn.Compiler
-  defdelegate __stream__(key, input, acc, vars, fun, args, opts), to: Nx.Defn.Evaluator
-
   @impl Nx.Defn.Compiler
   def __to_backend__(opts) do
     device = Keyword.get(opts, :device, :gpu)

lib/emlx/backend.ex

@@ -284,7 +284,7 @@ defmodule EMLX.Backend do
           {[i | axes], [max(low, 0) | lows], [max(high, 0) | highs]}
       end)
 
-    {_device, pad_value_mx} = from_nx(pad_value)
+    pad_value_mx = from_nx(pad_value)
 
     interior_padding = Enum.map(input_config, fn {_low, _high, interior} -> interior end)
 
@@ -1381,7 +1381,7 @@ defmodule EMLX.Backend do
     window_dilations = opts[:window_dilations] || List.duplicate(1, tuple_size(window_shape))
     interior_padding_config = Enum.map(window_dilations, &(&1 - 1))
 
-    {_device, zero_mx} = EMLX.scalar_tensor(0, :bool, device)
+    zero_mx = EMLX.scalar_tensor(0, :bool, device)
 
     window =
       1
@@ -1391,7 +1391,7 @@ defmodule EMLX.Backend do
 
     window_shape = EMLX.shape(window)
 
-    {device, pad_mx} =
+    pad_value_mx =
       case op do
         :sum ->
           EMLX.scalar_tensor(0, to_mlx_type(out.type), device)
@@ -1406,11 +1406,11 @@ defmodule EMLX.Backend do
           Nx.Constants.max(tensor.type, backend: {EMLX.Backend, device: device}) |> from_nx()
       end
 
-    padded_mx = EMLX.pad(t_mx, Nx.axes(tensor), low_pad, high_pad, pad_mx)
+    padded_mx = EMLX.pad(t_mx, Nx.axes(tensor), low_pad, high_pad, pad_value_mx)
 
     padded_mx
     |> sliding_window_view(EMLX.shape(padded_mx), window_shape, opts[:strides])
-    |> then(&EMLX.where(window, &1, {device, pad_mx}))
+    |> then(&EMLX.where(window, &1, pad_value_mx))
     |> then(&apply(EMLX, op, [&1, axes, false]))
     |> to_nx(out)
   end
@@ -1468,14 +1468,13 @@ defmodule EMLX.Backend do
   end
 
   defp window_scatter_function(function, out, tensor, source, init_value, window_dims_tuple, opts) do
-    # TODO: support window dilations
     unfold_flat = fn tensor ->
       {device, _} = t_mx = from_nx(tensor)
-      {_, pad_mx} = EMLX.scalar_tensor(0, EMLX.scalar_type(t_mx), device)
+      pad_value_mx = EMLX.scalar_tensor(0, EMLX.scalar_type(t_mx), device)
 
       {low_pad, high_pad} = Enum.unzip(opts[:padding])
 
-      padded_mx = EMLX.pad(t_mx, Nx.axes(tensor), low_pad, high_pad, pad_mx)
+      padded_mx = EMLX.pad(t_mx, Nx.axes(tensor), low_pad, high_pad, pad_value_mx)
 
       unfolded_mx =
         sliding_window_view(

mix.exs

@@ -58,7 +58,7 @@ defmodule EMLX.MixProject do
   defp deps do
     [
       {:elixir_make, "~> 0.6"},
-      {:nx, "~> 0.9.2"},
+      {:nx, "~> 0.10"},
       {:nif_call, "~> 0.1.3"}
     ]
   end

mix.lock

@@ -1,7 +1,7 @@
 %{
-  "complex": {:hex, :complex, "0.5.0", "af2d2331ff6170b61bb738695e481b27a66780e18763e066ee2cd863d0b1dd92", [:mix], [], "hexpm", "2683bd3c184466cfb94fad74cbfddfaa94b860e27ad4ca1bffe3bff169d91ef1"},
+  "complex": {:hex, :complex, "0.6.0", "b0130086a7a8c33574d293b2e0e250f4685580418eac52a5658a4bd148f3ccf1", [:mix], [], "hexpm", "0a5fa95580dcaf30fcd60fe1aaf24327c0fe401e98c24d892e172e79498269f9"},
   "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"},
   "nif_call": {:hex, :nif_call, "0.1.3", "bb4af0d28d1a2f10602d50246155b95b4ef6a389025c12d830dc9924ef06d324", [:mix], [], "hexpm", "48ba2e66c7d5aab4ca1a9656d3fa7326d317ba63bac64d24b22a05c8c8a9aff0"},
-  "nx": {:hex, :nx, "0.9.2", "17563029c01bf749aad3c31234326d7665abd0acc33ee2acbe531a4759f29a8a", [:mix], [{:complex, "~> 0.5", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "914d74741617d8103de8ab1f8c880353e555263e1c397b8a1109f79a3716557f"},
+  "nx": {:hex, :nx, "0.10.0", "128e4a094cb790f663e20e1334b127c1f2a4df54edfb8b13c22757ec33133b4f", [:mix], [{:complex, "~> 0.6", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "3db8892c124aeee091df0e6fbf8e5bf1b81f502eb0d4f5ba63e6378ebcae7da4"},
   "telemetry": {:hex, :telemetry, "1.3.0", "fedebbae410d715cf8e7062c96a1ef32ec22e764197f70cda73d82778d61e7a2", [:rebar3], [], "hexpm", "7015fc8919dbe63764f4b4b87a95b7c0996bd539e0d499be6ec9d7f3875b79e6"},
 }

test/emlx/nx_doctest_test.exs

@@ -48,10 +48,7 @@ defmodule EMLX.Nx.DoctestTest do
   @to_be_fixed [
     :moduledoc,
     # MLX sorts NaNs lowest, Nx sorts them highest
-    argsort: 2,
-    # Missing support for window dilations
-    window_scatter_max: 5,
-    window_scatter_min: 5
+    argsort: 2
   ]
 
   @not_supported [

test/emlx/nx_test.exs

@@ -1055,8 +1055,6 @@ defmodule EMLX.NxTest do
       )
     end
 
-    # does not support window_dilations yet
-    @tag :skip
     test "works with non-default options" do
       t = Nx.tensor([[[4, 2, 1, 3], [4, 2, 1, 7]], [[1, 2, 5, 7], [1, 8, 9, 2]]])
       opts = [strides: [2, 1, 1], padding: :valid, window_dilations: [1, 2, 2]]
@@ -1117,8 +1115,6 @@ defmodule EMLX.NxTest do
       )
     end
 
-    # window_dilations are not supported yet
-    @tag :skip
     test "works with non-default options" do
       t = Nx.tensor([[[4, 2, 1, 3], [4, 2, 1, 7]], [[1, 2, 5, 7], [1, 8, 9, 2]]])
       opts = [strides: [2, 1, 1], padding: :valid, window_dilations: [1, 2, 2]]
@@ -1244,8 +1240,6 @@ defmodule EMLX.NxTest do
       )
     end
 
-    # window dilations are not supported yet
-    @tag :skip
     test "supports window dilations" do
       result = Nx.window_sum(Nx.iota({4, 4}), {2, 2}, window_dilations: [2, 1])
 
@@ -1298,8 +1292,6 @@ defmodule EMLX.NxTest do
       )
     end
 
-    # window dilations are not supported yet
-    @tag :skip
     test "supports window dilations" do
       result = Nx.window_product(Nx.iota({4, 4}), {2, 2}, window_dilations: [2, 1])
 
@@ -1576,4 +1568,210 @@ defmodule EMLX.NxTest do
       )
     end
   end
+
+  describe "window_scatter_max" do
+    test "window_scatter_max with strides [2, 3]" do
+      t =
+        Nx.tensor([
+          [7, 2, 5, 3, 10, 2],
+          [3, 8, 9, 3, 4, 2],
+          [1, 5, 7, 5, 6, 1],
+          [0, 6, 2, 7, 2, 8]
+        ])
+
+      opts = [strides: [2, 3], padding: :valid]
+      result = Nx.window_scatter_max(t, Nx.tensor([[2, 6], [3, 1]]), 0, {2, 3}, opts)
+
+      assert_all_close(
+        result,
+        Nx.tensor([
+          [0, 0, 0, 0, 6, 0],
+          [0, 0, 2, 0, 0, 0],
+          [0, 0, 3, 0, 0, 0],
+          [0, 0, 0, 0, 0, 1]
+        ])
+      )
+    end
+
+    test "window_scatter_max with strides [2, 2]" do
+      t =
+        Nx.tensor([
+          [7, 2, 5, 3, 8],
+          [3, 8, 9, 3, 4],
+          [1, 5, 7, 5, 6],
+          [0, 6, 2, 10, 2]
+        ])
+
+      opts = [strides: [2, 2], padding: :valid]
+      result = Nx.window_scatter_max(t, Nx.tensor([[2, 6], [3, 1]]), 0, {2, 3}, opts)
+
+      assert_all_close(
+        result,
+        Nx.tensor([
+          [0, 0, 0, 0, 0],
+          [0, 0, 8, 0, 0],
+          [0, 0, 3, 0, 0],
+          [0, 0, 0, 1, 0]
+        ])
+      )
+    end
+
+    test "window_scatter_max with vectorized input" do
+      t =
+        Nx.tensor([
+          [
+            [7, 2, 5, 3],
+            [3, 8, 9, 3]
+          ],
+          [
+            [1, 5, 7, 5],
+            [0, 6, 2, 8]
+          ]
+        ])
+        |> Nx.vectorize(:x)
+
+      opts = [strides: [1, 2], padding: :valid]
+
+      source =
+        Nx.tensor([
+          [[2, 6]],
+          [[3, 1]]
+        ])
+        |> Nx.vectorize(:y)
+
+      result = Nx.window_scatter_max(t, source, 0, {2, 2}, opts)
+
+      Nx.Testing.assert_equal(
+        result,
+        Nx.tensor([
+          [
+            [
+              [0, 0, 0, 0],
+              [0, 2, 6, 0]
+            ],
+            [
+              [0, 0, 0, 0],
+              [0, 3, 1, 0]
+            ]
+          ],
+          [
+            [
+              [0, 0, 0, 0],
+              [0, 2, 0, 6]
+            ],
+            [
+              [0, 0, 0, 0],
+              [0, 3, 0, 1]
+            ]
+          ]
+        ])
+        |> Nx.vectorize([:x, :y])
+      )
+    end
+  end
+
+  describe "window_scatter_min" do
+    test "window_scatter_min with strides [2, 3]" do
+      t =
+        Nx.tensor([
+          [7, 2, 5, 3, 10, 2],
+          [3, 8, 9, 3, 4, 2],
+          [1, 5, 7, 5, 6, 1],
+          [0, 6, 2, 7, 2, 8]
+        ])
+
+      opts = [strides: [2, 3], padding: :valid]
+      result = Nx.window_scatter_min(t, Nx.tensor([[2, 6], [3, 1]]), 0, {2, 3}, opts)
+
+      assert_all_close(
+        result,
+        Nx.tensor([
+          [0, 2, 0, 0, 0, 0],
+          [0, 0, 0, 0, 0, 6],
+          [0, 0, 0, 0, 0, 1],
+          [3, 0, 0, 0, 0, 0]
+        ])
+      )
+    end
+
+    test "window_scatter_min with strides [2, 2]" do
+      t =
+        Nx.tensor([
+          [7, 2, 5, 3, 8],
+          [3, 8, 9, 3, 4],
+          [1, 5, 7, 5, 6],
+          [0, 6, 2, 10, 2]
+        ])
+
+      opts = [strides: [2, 2], padding: :valid]
+      result = Nx.window_scatter_min(t, Nx.tensor([[2, 6], [3, 1]]), 0, {2, 3}, opts)
+
+      assert_all_close(
+        result,
+        Nx.tensor([
+          [0, 2, 0, 0, 0],
+          [0, 0, 0, 6, 0],
+          [0, 0, 0, 0, 0],
+          [3, 0, 0, 0, 1]
+        ])
+      )
+    end
+
+    test "window_scatter_min with vectorized input" do
+      t =
+        Nx.tensor([
+          [
+            [7, 2, 5, 3],
+            [3, 8, 9, 3]
+          ],
+          [
+            [1, 5, 7, 5],
+            [0, 6, 2, 8]
+          ]
+        ])
+        |> Nx.vectorize(:x)
+
+      opts = [strides: [1, 2], padding: :valid]
+
+      source =
+        Nx.tensor([
+          [
+            [2, 6]
+          ],
+          [
+            [3, 1]
+          ]
+        ])
+        |> Nx.vectorize(:y)
+
+      result = Nx.window_scatter_min(t, source, 0, {2, 2}, opts)
+
+      assert_all_close(
+        result,
+        Nx.tensor([
+          [
+            [
+              [0, 2, 0, 0],
+              [0, 0, 0, 6]
+            ],
+            [
+              [0, 3, 0, 0],
+              [0, 0, 0, 1]
+            ]
+          ],
+          [
+            [
+              [0, 0, 0, 0],
+              [2, 0, 6, 0]
+            ],
+            [
+              [0, 0, 0, 0],
+              [3, 0, 1, 0]
+            ]
+          ]
+        ])
+        |> Nx.vectorize([:x, :y])
+      )
+    end
+  end
 end

test/support/emlx_case.ex

@@ -8,43 +8,7 @@ defmodule EMLX.Case do
   using do
     quote do
       import EMLX.Case
-    end
-  end
-
-  def assert_all_close(left, right, opts \\ []) do
-    atol = opts[:atol] || 1.0e-4
-    rtol = opts[:rtol] || 1.0e-4
-
-    equals =
-      left
-      |> Nx.all_close(right, atol: atol, rtol: rtol)
-      |> Nx.backend_transfer(Nx.BinaryBackend)
-
-    if equals != Nx.tensor(1, type: {:u, 8}, backend: Nx.BinaryBackend) do
-      flunk("""
-      Tensor assertion failed.
-      left: #{inspect(left)}
-      right: #{inspect(right)}
-      """)
-    end
-  end
-
-  def assert_equal(left, right) do
-    both_nan = Nx.is_nan(left) |> Nx.logical_and(Nx.is_nan(right))
-
-    equals =
-      left
-      |> Nx.equal(right)
-      |> Nx.logical_or(both_nan)
-      |> Nx.all()
-      |> Nx.to_number()
-
-    if equals != 1 do
-      flunk("""
-      Tensor assertion failed.
-      left: #{inspect(left)}
-      right: #{inspect(right)}
-      """)
+      import Nx.Testing
     end
   end
 end