feat: support tie breaks (#83)

* feat: support tie breaks

* chore: format

Author: polvalente

README.md

@@ -49,23 +49,6 @@ Defaulting to Nx.Defn.Evaluator is the safest option for now.
 Nx.Defn.default_options(compiler: EMLX)
 ```
 
-### Configuration
-
-EMLX supports several configuration options that can be set in your application's config:
-
-#### `:warn_unsupported_option`
-
-Controls whether warnings are logged when unsupported options are used with certain operations.
-
-- **Type**: `boolean`
-- **Default**: `true`
-- **Description**: When enabled, EMLX will log warnings for operations that receive options not supported by the MLX backend. For example, `Nx.argmax/2` and `Nx.argmin/2` with `tie_break: :high` will log a warning since MLX doesn't support this tie-breaking behavior.
-
-```elixir
-# In config/config.exs
-config :emlx, :warn_unsupported_option, false
-```
-
 ### MLX binaries
 
 EMLX relies on the [MLX](https://github.com/ml-explore/mlx) library to function, and currently EMLX will download precompiled builds from [mlx-build](https://github.com/cocoa-xu/mlx-build).

lib/emlx/backend.ex

@@ -246,7 +246,14 @@ defmodule EMLX.Backend do
 
   @impl true
   def reverse(out, tensor, axes) do
-    shape = Tuple.to_list(tensor.shape)
+    tensor
+    |> from_nx()
+    |> reverse_mlx(tensor.shape, axes)
+    |> to_nx(out)
+  end
+
+  defp reverse_mlx(tensor_mx, shape, axes) do
+    shape = Tuple.to_list(shape)
 
     {starts_stops, strides} =
       shape
@@ -264,10 +271,7 @@ defmodule EMLX.Backend do
 
     {starts, stops} = Enum.unzip(starts_stops)
 
-    tensor
-    |> from_nx()
-    |> EMLX.slice(starts, stops, strides)
-    |> to_nx(out)
+    EMLX.slice(tensor_mx, starts, stops, strides)
   end
 
   @impl true
@@ -584,21 +588,41 @@ defmodule EMLX.Backend do
       axis = opts[:axis]
       keep_axis = opts[:keep_axis] == true
 
-      if Application.get_env(:emlx, :warn_unsupported_option, true) and opts[:tie_break] == :high do
-        Logger.warning(
-          "Nx.Backend.#{unquote(op)}/3 with tie_break: :high is not supported in EMLX"
-        )
-      end
-
       t_mx = from_nx(tensor)
 
-      result =
+      t_mx =
+        if opts[:tie_break] == :high do
+          reverse_mlx(t_mx, tensor.shape, [axis] || Nx.axes(tensor))
+        else
+          t_mx
+        end
+
+      {device, _} =
+        result =
         if axis do
           EMLX.unquote(op)(t_mx, axis, keep_axis)
         else
           EMLX.unquote(op)(t_mx, keep_axis)
         end
 
+      # in case we had tie_break: :high, we need to subtract the result from the size of the sorted
+      # set because in reversing the axis above, we will get the complement of the result
+      result =
+        case {axis, opts[:tie_break]} do
+          {nil, :high} ->
+            size = EMLX.scalar_tensor(Nx.size(tensor) - 1, to_mlx_type(out.type), device)
+            EMLX.subtract(size, result)
+
+          {_, :high} ->
+            size =
+              EMLX.scalar_tensor(Nx.axis_size(tensor, axis) - 1, to_mlx_type(out.type), device)
+
+            EMLX.subtract(size, result)
+
+          {_, _} ->
+            result
+        end
+
       result
       |> EMLX.astype(to_mlx_type(out.type))
       |> to_nx(out)

test/emlx/config_test.exs

@@ -48,10 +48,8 @@ defmodule EMLX.Nx.DoctestTest do
   @to_be_fixed [
     :moduledoc,
     # MLX sorts NaNs lowest, Nx sorts them highest
-    argmin: 2,
-    argmax: 2,
     argsort: 2,
-    # Missing support for window dilations and for tie_break: :high
+    # Missing support for window dilations
     window_scatter_max: 5,
     window_scatter_min: 5
   ]
@@ -61,7 +59,10 @@ defmodule EMLX.Nx.DoctestTest do
     window_reduce: 5,
     population_count: 1,
     count_leading_zeros: 1,
-    sort: 2
+    sort: 2,
+    # We do not support the same ordering for NaNs as Nx
+    argmin: 2,
+    argmax: 2
   ]
 
   doctest Nx,