Allow passing a function to findmin and findmax (#673)

Author: kshyatt

src/host/indexing.jl

@@ -226,7 +226,7 @@ end
 Base.findfirst(A::AnyGPUArray{Bool}) = findfirst(identity, A)
 Base.findlast(A::AnyGPUArray{Bool})  = findlast(identity, A)
 
-function findminmax(binop, A::AnyGPUArray; init, dims)
+function findminmax(binop, f, A::AnyGPUArray; init, dims)
     indices = EachIndex(A)
     dummy_index = firstindex(A)
 
@@ -237,21 +237,25 @@ function findminmax(binop, A::AnyGPUArray; init, dims)
         isequal(x, y) && return (x, min(i, j))
         return t1
     end
+    
+    fA = f.(A)
 
     if dims == Colon()
-        res = mapreduce(tuple, reduction, A, indices; init = (init, dummy_index))
+        res = mapreduce(tuple, reduction, fA, indices; init = (init, dummy_index))
 
         # out of consistency with Base.findarray, return a CartesianIndex
         # when the input is a multidimensional array
         return (res[1], ndims(A) == 1 ? res[2] : CartesianIndices(A)[res[2]])
     else
-        res = mapreduce(tuple, reduction, A, indices;
+        res = mapreduce(tuple, reduction, fA, indices;
                         init = (init, dummy_index), dims=dims)
         vals = map(x->x[1], res)
         inds = map(x->ndims(A) == 1 ? x[2] : CartesianIndices(A)[x[2]], res)
         return (vals, inds)
     end
 end
 
-Base.findmax(a::AnyGPUArray; dims=:) = findminmax(Base.isless, a; init=typemin(eltype(a)), dims)
-Base.findmin(a::AnyGPUArray; dims=:) = findminmax(Base.isgreater, a; init=typemax(eltype(a)), dims)
+Base.findmax(a::AnyGPUArray; dims=:) = findminmax(Base.isless, identity, a; init=typemin(eltype(a)), dims)
+Base.findmin(a::AnyGPUArray; dims=:) = findminmax(Base.isgreater, identity, a; init=typemax(eltype(a)), dims)
+Base.findmax(f::Function, a::AnyGPUArray; dims=:) = findminmax(Base.isless, f, a; init=typemin(eltype(a)), dims)
+Base.findmin(f::Function, a::AnyGPUArray; dims=:) = findminmax(Base.isgreater, f, a; init=typemax(eltype(a)), dims)

test/testsuite/indexing.jl

@@ -205,6 +205,14 @@ end
             @test isequal(findmax(x), findmax(AT(x)))
             @test isequal(findmax(x; dims=1), Array.(findmax(AT(x); dims=1)))
         end
+        let x = randn(Float32, 100)
+            @test findmax(abs, x) == findmax(abs, AT(x))
+            @test findmax(abs, x; dims=1) == Array.(findmax(abs, AT(x); dims=1))
+
+            x[32] = x[33] = x[55] = x[66] = NaN32
+            @test isequal(findmax(abs, x), findmax(abs, AT(x)))
+            @test isequal(findmax(abs, x; dims=1), Array.(findmax(abs, AT(x); dims=1)))
+        end
         let x = rand(Float32, 10, 10)
             @test findmax(x) == findmax(AT(x))
             @test findmax(x; dims=1) == Array.(findmax(AT(x); dims=1))
@@ -235,6 +243,14 @@ end
             @test isequal(findmin(x), findmin(AT(x)))
             @test isequal(findmin(x; dims=1), Array.(findmin(AT(x); dims=1)))
         end
+        let x = randn(Float32, 100)
+            @test findmin(abs, x) == findmin(abs, AT(x))
+            @test findmin(abs, x; dims=1) == Array.(findmin(abs, AT(x); dims=1))
+
+            x[32] = x[33] = x[55] = x[66] = NaN32
+            @test isequal(findmin(abs, x), findmin(abs, AT(x)))
+            @test isequal(findmin(abs, x; dims=1), Array.(findmin(abs, AT(x); dims=1)))
+        end
         let x = rand(Float32, 10, 10)
             @test findmin(x) == findmin(AT(x))
             @test findmin(x; dims=1) == Array.(findmin(AT(x); dims=1))