feat: add Nx.argsort/2

[polvalente]

This PR aims to add the sort_indices option so we can support "argsort" functionality

closes #342

[polvalente]

This isn't ready because EXLA doesn't implement the option yet.
@seanmor5 any pointers on how we can achieve this?

I read the XLA docs on sort (https://www.tensorflow.org/xla/operation_semantics#sort) and it seems that if we provide 2 arguments, they will be sorted simultaneously. However, I had some trouble on how to build the corresponding index tensor inside the NIF.

I'm leaning towards using something akin to the code below for build the second argument:

iex(7)> tensor = Nx.tensor([[[[4,3,2],[1,0,-1]], [[1,2,3],[4,5,2]]]])
#Nx.Tensor<
  s64[1][2][2][3]
  [
    [
      [
        [4, 3, 2],
        [1, 0, -1]
      ],
      [
        [1, 2, 3],
        [4, 5, 2]
      ]
    ]
  ]
>
iex(8)> Nx.iota(tensor, axis: Nx.axes(tensor) |> Enum.at(-1))        
#Nx.Tensor<
  s64[1][2][2][3]
  [
    [
      [
        [0, 1, 2],
        [0, 1, 2]
      ],
      [
        [0, 1, 2],
        [0, 1, 2]
      ]
    ]
  ]
>

However, I'm not sure how it will behave, if it'll actually need a arity-4 function for sorting (as implied by the documentation linked above).

[josevalim]

Awesome job @polvalente!

Some notes:

1. Generally speaking, we prefer to expose the EXLA operations as close to "native" as possible. I.e. EXLA.Op should expose the C API instead of providing conveniences on top of it. For this purpose, it is probably better to implement variadic_sort, pretty much like we implemented variadic_reduce.

2. At the NX level, I would rather introduce an argsort function instead of having options inside sort. :)

[seanmor5]

@polvalente I looked into implementing argsort in EXLA a little bit ago. What I would do is change the sort NIF to always accept a list of Tensors. Then you can implement argsort in EXLA.Lib and sort the input tensor side by side with an iota operation.

EDIT: Actually I like the variadic sort idea better, I would go with that :)

[polvalente]

I'm not sure if there's a better way to achieve this. Basically, the comparison function needed for variadic_sort always has 2*num_args arity, so I needed to extend the arity-2 comparator to an arity-4 one which only uses the first 2 args (which are the current comparison elements for the first of the args passed to sort)

[josevalim]

I think you can wrap that function inside the existing function. Something like this:

    # Grow the comparator to arity 4 because argsort uses
    # variadic_sort underneath
    [[arg1, arg2], _expr, fun] = comparator.data.args

    comparator_4 = Expr.fun([arg1, arg2, arg1, arg2], fn x, y, _, _ -> fun.(x, y) end)

[polvalente]

I wasn't able to make argsort work with {:u, 64} type on EXLA, which would be ideal, so I defaulted to {:s, 64}

[josevalim]

Which error do you get? Looking at the code it feels like it should work... on the other, numpy and jax both return int64 too.

[josevalim]

Torchx also return int64, so let's keep this as is.

[polvalente]

The error was something along the lines of "The comparator function expected 0-th argument to have shape s64[], got u64[] instead"

[polvalente]

These changes are two-fold. First because the original comparators didn't work with negative numbers, and secondly I've used<= and >= so the sort is stable

[polvalente]

I'm not sure if the match_types call on line 1654 is needed. Perhaps we should remove it or refactor the for-loop somehow to use it correctly

[josevalim]

It isn't, you can remove it.

[josevalim]

The default here seems to be :desc. We probably want to change both sort and argsort to be asc.

[polvalente]

It seems that the comparator definitions were swapped between desc and asc. My latest commit changed this

[josevalim]

Two minor comments and ship it! 🚢