test: more tests broken

Author: avik-pal

test/ops.jl

@@ -1227,7 +1227,8 @@ end
         x_ra = Reactant.to_rarray(randn(Float32, 6, 6))
         lu_ra, ipiv, perm, info = @jit Ops.lu(x_ra)
 
-        @test @jit(recon_from_lu(lu_ra)) ≈ @jit(getindex(x_ra, perm, :))
+        @test @jit(recon_from_lu(lu_ra)) ≈ @jit(getindex(x_ra, perm, :)) atol = 1e-5 rtol =
+            1e-2
     end
 
     @testset "batched" begin

test/sorting.jl

@@ -176,22 +176,38 @@ end
 end
 
 @testset "findfirst / findlast" begin
-    x = rand(Bool, 3, 4)
+    x = Bool[
+        0 0 0 0
+        1 0 1 0
+        0 1 0 1
+    ]
     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))
+    @test ffirstlinindices(x) ≈ @jit(findfirst(x_ra)) broken = contains(
+        string(Reactant.devices()[1]), "TPU"
+    )
+    @test flastlinindices(x) ≈ @jit(findlast(x_ra)) broken = contains(
+        string(Reactant.devices()[1]), "TPU"
+    )
+
+    x = Int64[
+        3 5 7 9
+        4 6 7 8
+        5 7 8 9
+    ]
+    x_ra = Reactant.to_rarray(x)
+
+    @test ffirstlinindices(iseven, x) ≈ @jit(findfirst(iseven, x_ra)) broken = contains(
+        string(Reactant.devices()[1]), "TPU"
+    )
+    @test flastlinindices(iseven, x) ≈ @jit(findlast(iseven, x_ra)) broken = contains(
+        string(Reactant.devices()[1]), "TPU"
+    )
 end
 
 @testset "approx top k lowering" begin