feature request: support argsort/argmax-style interface

[tmbdev]

Julia currently supports 1D sortperm, and PR #45211 adds sortperm support for multidmensional arrays, returning CartesianIndex.

PyTorch and NumPy both provide multidimensional argsort functions that return sort indexes along only one dimension. While this is less general, it requires less storage in the return value for high dimensional arrays. Furthermore, there are CUDA kernels available for argsort but not for returning CartesianIndex. A similar comment applies to findmax vs NumPy argmax.

I would find it very useful if Julia could support argsort/argmax-style interfaces in addition to sortperm/findmax-style interfaces.

[LilithHafner]

Thank you for noticing this and opening an issue. I do think we need to reconsider the return values of our perm/arg/find functions.

PR #45211 adds sortperm support for multidmensional arrays, returning CartesianIndex.

We do not use CartesianIndex here, so memory footprint should not be an issue:

julia> x = rand(3, 3, 3, 3);

julia> y = sortperm(x; dims=3)
3×3×3×3 Array{Int64, 4}:
...

julia> Base.summarysize(x)
704

julia> Base.summarysize(y)
704

[tmbdev]

But CartesianIndex is returned for argmin:

julia> size(ds)
(2048, 512)

julia> argmin(ds, dims=1)
1×512 CuArray{CartesianIndex{2}, 2, CUDA.Mem.DeviceBuffer}:
 CartesianIndex(1271, 1)  CartesianIndex(9, 2)  CartesianIndex(9, 3)  CartesianIndex(9, 4)  CartesianIndex(9, 5)  …  CartesianIndex(250, 509)  CartesianIndex(2014, 510)  CartesianIndex(9, 511)  CartesianIndex(9, 512)

julia> 

I think the consistent thing to do would be for argmin/argmax/sortperm to always return CartesianIndex, and to have a separate set of functions argmin1/argmax1/sortperm1 that always only operate over one dimension and consistently return integer indexes referring to slices along that dimension.

[LilithHafner]

Perhaps you are looking for argmin.(eachslice(ds, dims=2))?

[tmbdev]

eachslice is much slower on the GPU.

using CUDA
CUDA.allowscalar(false)

a = randn(10000, 10000)
ca = CuArray(a)

for _ in 1:3
    println()
    @time argmin(a, dims=2);
    @time argmin.(eachslice(a, dims=2));
    @time argmin(ca, dims=2);
    @time argmin.(eachslice(ca, dims=2));
end

  0.242818 seconds (9 allocations: 234.562 KiB)
  0.201610 seconds (9 allocations: 468.938 KiB)
  0.000097 seconds (58 allocations: 2.438 KiB)
  0.311594 seconds (420.59 k allocations: 21.550 MiB)

The issue isn't whether one can somehow compute argmin, it's whether one can compute it efficiently and with minimal overhead on both CPU and GPU.

[LilithHafner]

That sounds like a feature request for CUDA.jl.

[tmbdev]

That sounds complicated; is there even an API to fuse "op.(broadcast(...))" into a single kernel? Even if there were, you end up with an interface that most people coming to Julia simply won't think of or have any reason to assume that it's efficient. I think you're making something pretty straightforward rather complicated.

My suggestion remains therefore for Julia to support the same kind of simple argmin/argmax/argsort/minloc API also found in Fortran, NumPy, PyTorch, both on CPU and GPU. The GPU kernels already exist, it's what people are familiar with, it's efficient, and it would make porting code easier.