fix: more tests and fixes for find functions

Author: avik-pal

src/TracedRArray.jl

@@ -836,14 +836,14 @@ Base.findfirst(x::AnyTracedRArray) = findfirst(identity, x)
 Base.findlast(x::AnyTracedRArray) = findlast(identity, x)
 
 function Base.findfirst(f::Function, x::AnyTracedRArray)
-    fA = f.(x)
-    (; indices) = Ops.top_k(materialize_traced_array(fA), 1)
+    fA = materialize_traced_array(vec(f.(x)))
+    (; indices) = Ops.top_k(fA, 1)
     return @allowscalar indices[1]
 end
 
 function Base.findlast(f::Function, x::AnyTracedRArray)
     fA = Ops.reverse(materialize_traced_array(vec(f.(x))); dimensions=[1])
-    (; indices) = Ops.top_k(materialize_traced_array(fA), 1)
+    (; indices) = Ops.top_k(fA, 1)
     return length(x) - @allowscalar(indices[1]) + 1
 end
 
@@ -883,7 +883,9 @@ function Base.findmin(f, x::AnyTracedRArray; dims::Union{Integer,Nothing}=nothin
         )
     end
 
-    return (Ops.negate(values), linear_indices)
+    values = Ops.negate(values)
+    ndims(x) == 1 && return @allowscalar (values[1], linear_indices[1])
+    return (values, linear_indices)
 end
 
 function Base.findmax(f, x::AnyTracedRArray; dims::Union{Integer,Nothing}=nothing)
@@ -910,6 +912,7 @@ function Base.findmax(f, x::AnyTracedRArray; dims::Union{Integer,Nothing}=nothin
         )
     end
 
+    ndims(x) == 1 && return @allowscalar (values[1], linear_indices[1])
     return (values, linear_indices)
 end
 

test/sorting.jl

@@ -39,3 +39,54 @@ using Reactant, Test
     @test argmin(abs2, x) == @jit(argmin(abs2, x_ra))
     @test argmax(abs2, x) == @jit(argmax(abs2, x_ra))
 end
+
+@testset "findmin / findmax" begin
+    xvec = randn(10)
+    xvec_ra = Reactant.to_rarray(xvec)
+
+    x = randn(2, 3)
+    x_ra = Reactant.to_rarray(x)
+
+    function fwithlinindices(g, f, x; kwargs...)
+        values, indices = g(f, x; kwargs...)
+        return values, LinearIndices(x)[indices]
+    end
+
+    @test fwithlinindices(findmin, identity, x) == @jit(findmin(x_ra))
+    @test fwithlinindices(findmax, identity, x) == @jit(findmax(x_ra))
+    @test fwithlinindices(findmin, identity, xvec) == @jit(findmin(xvec_ra))
+    @test fwithlinindices(findmax, identity, xvec) == @jit(findmax(xvec_ra))
+
+    fmindims(x, d) = findmin(x; dims=d)
+    fmindims(f, x, d) = findmin(f, x; dims=d)
+    fmaxdims(x, d) = findmax(x; dims=d)
+    fmaxdims(f, x, d) = findmax(f, x; dims=d)
+
+    @test fwithlinindices(findmin, identity, x; dims=1) == @jit(fmindims(x_ra, 1))
+    @test fwithlinindices(findmax, identity, x; dims=1) == @jit(fmaxdims(x_ra, 1))
+    @test fwithlinindices(findmin, identity, x; dims=2) == @jit(fmindims(x_ra, 2))
+    @test fwithlinindices(findmax, identity, x; dims=2) == @jit(fmaxdims(x_ra, 2))
+    @test fwithlinindices(findmin, abs2, x; dims=1) == @jit(fmindims(abs2, x_ra, 1))
+    @test fwithlinindices(findmax, abs2, x; dims=1) == @jit(fmaxdims(abs2, x_ra, 1))
+    @test fwithlinindices(findmin, abs2, x; dims=2) == @jit(fmindims(abs2, x_ra, 2))
+    @test fwithlinindices(findmax, abs2, x; dims=2) == @jit(fmaxdims(abs2, x_ra, 2))
+end
+
+@testset "findfirst / findlast" begin
+    x = rand(Bool, 3, 4)
+    x_ra = Reactant.to_rarray(x)
+
+    ffirstlinindices(x) = LinearIndices(x)[findfirst(x)]
+    ffirstlinindices(f, x) = LinearIndices(x)[findfirst(f, x)]
+    flastlinindices(x) = LinearIndices(x)[findlast(x)]
+    flastlinindices(f, x) = LinearIndices(x)[findlast(f, x)]
+
+    @test ffirstlinindices(x) == @jit(findfirst(x_ra))
+    @test flastlinindices(x) == @jit(findlast(x_ra))
+
+    x = rand(1:256, 3, 4)
+    x_ra = Reactant.to_rarray(x)
+
+    @test ffirstlinindices(iseven, x) == @jit(findfirst(iseven, x_ra))
+    @test flastlinindices(iseven, x) == @jit(findlast(iseven, x_ra))
+end