more tpu fixes

Author: wsmoses

test/basic.jl

@@ -988,13 +988,15 @@ end
     @test Array(x) ≈ Array(y) ./ 2
 end
 
+if !contains(string(Reactant.devices()[1]), "TPU")
 @testset "Hlo Cost Analysis" begin
     x_ra = Reactant.to_rarray(rand(4, 4))
     mul_comp = @compile x_ra * x_ra
     cost = Reactant.XLA.cost_analysis(mul_comp)
 
     @test cost isa Reactant.XLA.HloCostAnalysisProperties
 end
+end
 
 function fractional_idx(times, t)
     n₂ = searchsortedfirst(times, t)

test/constructor.jl

@@ -105,23 +105,23 @@ end
     rg = Reactant.to_rarray(g)
 
     @jit update!(rg)
-    @test convert(Array, rg.data) == [2.7, 1.59]
+    @test convert(Array, rg.data) ≈ [2.7, 1.59]
 
     rg = Reactant.to_rarray(g)
     res = @jit selfreturn(rg)
-    @test convert(Array, res.data) == [3.14, 1.59]
-    @test res.radius == 2.7
+    @test convert(Array, res.data) ≈ [3.14, 1.59]
+    @test res.radius ≈ 2.7
     @test typeof(res.radius) <: ConcreteRNumber
 
     rg = Reactant.to_rarray(g)
 
     @jit call_update!(rg)
-    @test convert(Array, rg.data) == [2.7, 1.59]
+    @test convert(Array, rg.data) ≈ [2.7, 1.59]
 
     rg = Reactant.to_rarray(g)
     res = @jit call_selfreturn(rg)
-    @test convert(Array, res.data) == [3.14, 1.59]
-    @test res.radius == 2.7
+    @test convert(Array, res.data) ≈ [3.14, 1.59]
+    @test res.radius ≈ 2.7
     @test typeof(res.radius) <: ConcreteRNumber
 end
 
@@ -131,7 +131,7 @@ end
 
     rg = Reactant.to_rarray(g)
     res = @jit selfreturn(rg)
-    @test convert(Array, res[1][].data) == [3.14, 1.59]
+    @test convert(Array, res[1][].data) ≈ [3.14, 1.59]
     @test convert(Array, res[2][].data) ≈ [3.14, 1.59]
     @test res[1][].data == res[2][].data
 end

test/ops.jl

@@ -10,6 +10,7 @@ using SpecialFunctions: SpecialFunctions
     x = Reactant.to_rarray([1.0, -1.0])
     @test [1.0, 1.0] ≈ @jit Ops.abs(x)
 
+    if !contains(string(Reactant.devices()[1]), "TPU")
     x = Reactant.to_rarray([
         3.0+4im -3.0+4im
         3.0-4im -3.0-4im
@@ -18,6 +19,7 @@ using SpecialFunctions: SpecialFunctions
         5.0 5.0
         5.0 5.0
     ] ≈ @jit Ops.abs(x)
+    end
 end
 
 @testset "add" begin
@@ -95,6 +97,7 @@ end
     @test cholesky(Array(x)).U ≈ @jit g1(x)
     @test transpose(cholesky(Array(x)).U) ≈ @jit g2(x)
 
+    if !contains(string(Reactant.devices()[1]), "TPU")
     x = Reactant.to_rarray(
         [
             10.0+0.0im 2.0-3.0im 3.0-4.0im
@@ -105,6 +108,7 @@ end
 
     @test cholesky(Array(x)).U ≈ @jit g1(x)
     @test adjoint(cholesky(Array(x)).U) ≈ @jit g2(x)
+   end
 end
 
 @testset "clamp" begin
@@ -135,11 +139,12 @@ end
         ),
     ]
         x = Reactant.to_rarray([1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8, 9.9, 10.0])
-        @test [3.0, 3.0, 3.3, 4.4, 5.5, 6.6, 7.0, 7.0, 7.0, 7.0] ==
+        @test [3.0, 3.0, 3.3, 4.4, 5.5, 6.6, 7.0, 7.0, 7.0, 7.0] ≈
             @jit Ops.clamp(_min, x, _max)
     end
 end
 
+if !contains(string(Reactant.devices()[1]), "TPU")
 @testset "complex" begin
     x = Reactant.to_rarray(1.1; track_numbers=true)
     y = Reactant.to_rarray(2.2; track_numbers=true)
@@ -149,6 +154,7 @@ end
     y = Reactant.to_rarray([5.5, 6.6, -7.7, -8.8])
     @test [1.1 + 5.5im, 2.2 + 6.6im, 3.3 - 7.7im, 4.4 - 8.8im] ≈ @jit Ops.complex(x, y)
 end
+end
 
 @testset "constant" begin
     for x in [[1, 2, 3], [1.1, 2.2, 3.3], [1.1 + 2.2im, 3.3 + 4.4im, 5.5 + 6.6im]]