more testing

Author: wsmoses

test/basic.jl

@@ -1027,7 +1027,7 @@ end
     times = 0:0.01:4.5
     @test times isa Base.StepRangeLen
     res = @jit fractional_idx(times, ConcreteRNumber(2.143))
-    @test res[1] == 0.29999999999997334
+    @test res[1] ≈ 0.29999999999997334
     @test res[2] == 215
     @test res[3] == 216
 end
@@ -1036,7 +1036,7 @@ end
     times = Reactant.to_rarray(0:0.01:4.5; track_numbers=Number)
     @test times isa Reactant.TracedRNumberOverrides.TracedStepRangeLen
     res = @jit fractional_idx(times, ConcreteRNumber(2.143))
-    @test res[1] == 0.29999999999997334
+    @test res[1] ≈ 0.29999999999997334
     @test res[2] == 215
     @test res[3] == 216
 end
@@ -1224,7 +1224,6 @@ end
 end
 
 accum_fn(x, y) = abs2(x) + abs2(y)
-Base.reduce_empty(::typeof(accum_fn), ::Type{T}) where {T} = zero(T)
 
 @testset "accumulate" begin
     a = collect(Float32, 1:10) ./ 10
@@ -1292,6 +1291,7 @@ Base.reduce_empty(::typeof(accum_fn), ::Type{T}) where {T} = zero(T)
         end ≈ cumprod(b; dims=3)
     end
 
+    if !contains(string(Reactant.devices()[1]), "TPU")
     @testset "accumulate" begin
         @test @jit(accumulate(accum_fn, a_ra; init=0.0f0)) ≈
             accumulate(accum_fn, a; init=0.0f0)
@@ -1325,6 +1325,7 @@ Base.reduce_empty(::typeof(accum_fn), ::Type{T}) where {T} = zero(T)
             z
         end ≈ accumulate(accum_fn, b; dims=3, init=0.0f0)
     end
+    end
 end
 
 sameunitrange(x, y) = first(x) == first(y) && last(x) == last(y)

test/ops.jl

@@ -1238,7 +1238,7 @@ end
         @test size(perm) == (4, 3, 6)
         @test size(info) == (4, 3)
 
-        @test @jit(recon_from_lu(lu_ra)) ≈ @jit(apply_permutation(x_ra, perm))
+        @test @jit(recon_from_lu(lu_ra)) ≈ @jit(apply_permutation(x_ra, perm)) atol = 1e-5 rtol = 1e-2
     end
 end
 

test/sorting.jl

@@ -134,8 +134,8 @@ end
     x = randn(2, 3, 4)
     x_ra = Reactant.to_rarray(x)
 
-    @test argmin(abs2, x) == @jit(argmin(abs2, x_ra))
-    @test argmax(abs2, x) == @jit(argmax(abs2, x_ra))
+    @test argmin(abs2, x) ≈ @jit(argmin(abs2, x_ra))
+    @test argmax(abs2, x) ≈ @jit(argmax(abs2, x_ra))
 end
 
 @testset "findmin / findmax" begin
@@ -152,22 +152,22 @@ 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))
+    @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))
+    @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
@@ -179,14 +179,14 @@ end
     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))
+    @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))
+    @test ffirstlinindices(iseven, x) ≈ @jit(findfirst(iseven, x_ra))
+    @test flastlinindices(iseven, x) ≈ @jit(findlast(iseven, x_ra))
 end
 
 @testset "approx top k lowering" begin

test/wrapped_arrays.jl

@@ -241,7 +241,7 @@ end
 @testset "reshaped subarray indexing" begin
     fn(x) = view(x, 1:2) .+ 1
     x_ra = Reactant.to_rarray(rand(3, 4, 3))
-    @test @jit(fn(x_ra)) == fn(Array(x_ra))
+    @test @jit(fn(x_ra)) ≈ fn(Array(x_ra))
 end
 
 function reshape_getindex(x)