test: more test fixes

Author: avik-pal

test/autodiff.jl

@@ -255,7 +255,7 @@ end
 
 @testset "seed" begin
     x = Reactant.to_rarray(rand(2, 2))
-    st = (; rng=Reactant.ConcreteRNG())
+    st = (; rng=Reactant.ReactantRNG())
 
     @test begin
         hlo = @code_hlo gradient_fn(x, st)

test/basic.jl

@@ -416,9 +416,9 @@ end
 
 @testset "Complex runtime: $CT" for CT in (ComplexF32, ComplexF64)
     # complex f64 not supported on tpu
-    a = Reactant.to_rarray(ones(CT, 2))
-    b = Reactant.to_rarray(ones(CT, 2))
     @test begin
+        a = Reactant.to_rarray(ones(CT, 2))
+        b = Reactant.to_rarray(ones(CT, 2))
         c = Reactant.compile(+, (a, b))(a, b)
         c == ones(CT, 2) + ones(CT, 2)
     end broken = CT != ComplexF32 && RunningOnTPU
@@ -987,11 +987,12 @@ end
     @test Array(x) ≈ Array(y) ./ 2
 end
 
-@test "HLO Cost Analysis" begin
+@testset "HLO Cost Analysis" begin
     x_ra = Reactant.to_rarray(rand(4, 4))
     mul_comp = @compile x_ra * x_ra
-    @test Reactant.XLA.cost_analysis(mul_comp) isa Reactant.XLA.HloCostAnalysisProperties broken =
-        RunningOnTPU
+    @test begin
+        Reactant.XLA.cost_analysis(mul_comp) isa Reactant.XLA.HloCostAnalysisProperties
+    end broken = RunningOnTPU
 end
 
 function fractional_idx(times, t)

test/indexing.jl

@@ -1,5 +1,7 @@
 using LinearAlgebra, Reactant, Test
 
+const RunningOnTPU = contains(string(Reactant.devices()[1]), "TPU")
+
 function update_on_copy(x)
     y = x[1:2, 2:4, :]
     y[1:1, 1:1, :] = ones(1, 1, 3)

test/ops.jl

@@ -17,10 +17,12 @@ const RunningOnAppleX86 = Sys.isapple() && Sys.ARCH === :x86_64
         3.0+4im -3.0+4im
         3.0-4im -3.0-4im
     ])
-    @test [
-        5.0 5.0
-        5.0 5.0
-    ] ≈ @jit Ops.abs(x) broken = RunningOnTPU
+    @test begin
+        [
+            5.0 5.0
+            5.0 5.0
+        ] ≈ @jit(Ops.abs(x))
+    end broken = RunningOnTPU
 end
 
 @testset "add" begin
@@ -40,7 +42,7 @@ end
     b = Reactant.to_rarray([
         9.9 + 10.10im, 11.11 + 12.12im, -13.13 + -14.14im, -15.15 + -16.16im
     ])
-    @test Array(a) .+ Array(b) ≈ @jit Ops.add(a, b) broken = RunningOnTPU
+    @test Array(a) .+ Array(b) ≈ @jit(Ops.add(a, b)) broken = RunningOnTPU
 end
 
 @testset "after_all" begin