test: more tests work

Author: avik-pal

test/autodiff.jl

@@ -196,17 +196,17 @@ end
 end
 
 @testset "Seed initialization of Complex arrays on matmul: Issue #593" begin
-    a = ones(ComplexF64, 2, 2)
-    b = 2.0 * ones(ComplexF64, 2, 2)
-    a_re = Reactant.to_rarray(a)
-    b_re = Reactant.to_rarray(b)
     df(x, y) = Enzyme.gradient(ReverseWithPrimal, *, x, y)
     @test begin
+        a = ones(ComplexF64, 2, 2)
+        b = 2.0 * ones(ComplexF64, 2, 2)
+        a_re = Reactant.to_rarray(a)
+        b_re = Reactant.to_rarray(b)
         res = @jit df(a_re, b_re) # before, this segfaulted
         (res.val ≈ 4ones(2, 2)) &&
             (res.derivs[1] ≈ 4ones(2, 2)) &&
             (res.derivs[2] ≈ 2ones(2, 2))
-    end broken = contains(string(Reactant.devices()[1]), "TPU")
+    end skip = contains(string(Reactant.devices()[1]), "TPU")
 end
 
 @testset "onehot" begin

test/basic.jl

@@ -414,16 +414,6 @@ end
     @test eltype(f(y)) == eltype(x)
 end
 
-@testset "Complex runtime: $CT" for CT in (ComplexF32, ComplexF64)
-    # complex f64 not supported on tpu
-    @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
-end
-
 @testset "Scalars" begin
     @testset "Only Scalars" begin
         x = (3, 3.14)
@@ -783,7 +773,7 @@ end
     @test begin
         x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN] .* im)
         @jit(isfinite.(x)) == [true, false, false, false, false]
-    end broken = RunningOnTPU
+    end skip = RunningOnTPU
 end
 
 @testset "isnan" begin
@@ -793,7 +783,7 @@ end
     @test begin
         x = Reactant.to_rarray([1.0, NaN, Inf, -Inf, NaN] .* im)
         @jit(isnan.(x)) == [false, true, false, false, true]
-    end broken = RunningOnTPU
+    end skip = RunningOnTPU
 end
 
 @testset "isnan/isfinite" begin

test/complex.jl

@@ -3,120 +3,82 @@ using Reactant
 
 const RunningOnTPU = contains(string(Reactant.devices()[1]), "TPU")
 
-@testset "conj" begin
-    @testset "$(typeof(x))" for x in (1.0, 1.0 + 2.0im)
-        x_concrete = Reactant.to_rarray(x)
-        @test only(@jit(conj(x_concrete))) == conj(x) broken = RunningOnTPU
-    end
-
-    @testset "$(typeof(x))" for x in (
-        fill(1.0 + 2.0im),
-        fill(1.0),
-        [1.0 + 2.0im; 3.0 + 4.0im],
-        [1.0; 3.0],
-        [1.0 + 2.0im 3.0 + 4.0im],
-        [1.0 2.0],
-        [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
-        [1.0 3.0; 5.0 7.0],
-    )
-        x_concrete = Reactant.to_rarray(x)
-        @test @jit(conj(x_concrete)) == conj(x) broken = RunningOnTPU
-    end
+@testset "Complex runtime: $CT" for CT in (ComplexF32, ComplexF64)
+    @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 skip = CT == ComplexF64 && RunningOnTPU
 end
 
-@testset "conj!" begin
-    @testset "$(typeof(x))" for x in (
-        fill(1.0 + 2.0im),
-        fill(1.0),
-        [1.0 + 2.0im; 3.0 + 4.0im],
-        [1.0; 3.0],
-        [1.0 + 2.0im 3.0 + 4.0im],
-        [1.0 2.0],
-        [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
-        [1.0 3.0; 5.0 7.0],
-    )
-        x_concrete = Reactant.to_rarray(x)
-        @test @jit(conj!(x_concrete)) == conj(x) broken = RunningOnTPU
-        @test x_concrete == conj(x) broken = RunningOnTPU
-    end
-end
+const SCALAR_LIST = (1.0, 1.0 + 2.0im)
 
-@testset "real" begin
-    @testset "$(typeof(x))" for x in (1.0, 1.0 + 2.0im)
-        x_concrete = Reactant.to_rarray(x)
-        @test only(@jit(real(x_concrete))) == real(x) broken = RunningOnTPU
-    end
-
-    @testset "$(typeof(x))" for x in (
-        fill(1.0 + 2.0im),
-        fill(1.0),
-        [1.0 + 2.0im; 3.0 + 4.0im],
-        [1.0; 3.0],
-        [1.0 + 2.0im 3.0 + 4.0im],
-        [1.0 2.0],
-        [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
-        [1.0 3.0; 5.0 7.0],
-    )
-        x_concrete = Reactant.to_rarray(x)
-        @test @jit(real(x_concrete)) == real(x) broken = RunningOnTPU
-    end
-end
+const ARRAY_LIST = (
+    fill(1.0 + 2.0im),
+    fill(1.0),
+    [1.0 + 2.0im; 3.0 + 4.0im],
+    [1.0; 3.0],
+    [1.0 + 2.0im 3.0 + 4.0im],
+    [1.0 2.0],
+    [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
+    [1.0 3.0; 5.0 7.0],
+)
 
-@testset "imag" begin
-    @testset "$(typeof(x))" for x in (1.0, 1.0 + 2.0im)
-        x_concrete = Reactant.to_rarray(x)
-        @test only(@jit(imag(x_concrete))) == imag(x)
+@testset "$(string(fn))" for fn in (conj, conj!, real, imag)
+    if !endswith(string(fn), "!")
+        @testset "$(typeof(x))" for x in SCALAR_LIST
+            @test begin
+                x_concrete = Reactant.to_rarray(x)
+                only(@jit(fn(x_concrete))) == fn(x)
+            end skip = RunningOnTPU && eltype(x) == ComplexF64
+        end
     end
 
-    @testset "$(typeof(x))" for x in (
-        fill(1.0 + 2.0im),
-        fill(1.0),
-        [1.0 + 2.0im; 3.0 + 4.0im],
-        [1.0; 3.0],
-        [1.0 + 2.0im 3.0 + 4.0im],
-        [1.0 2.0],
-        [1.0+2.0im 3.0+4.0im; 5.0+6.0im 7.0+8.0im],
-        [1.0 3.0; 5.0 7.0],
-    )
-        x_concrete = Reactant.to_rarray(x)
-        @test @jit(imag(x_concrete)) == imag(x) broken = RunningOnTPU
+    @testset "$(typeof(x))" for x in ARRAY_LIST
+        @test begin
+            x_concrete = Reactant.to_rarray(x)
+            @jit(fn(x_concrete)) == fn(x)
+        end skip = RunningOnTPU && eltype(x) == ComplexF64
     end
 end
 
 @testset "abs: $T" for T in (Float32, ComplexF32)
     x = randn(T, 10)
     x_concrete = Reactant.to_rarray(x)
-    @test @jit(abs.(x_concrete)) ≈ abs.(x) broken = RunningOnTPU
+    @test @jit(abs.(x_concrete)) ≈ abs.(x)
 end
 
 @testset "promote_to Complex" begin
-    x = 1.0 + 2.0im
+    x = ComplexF32(1.0 + 2.0im)
     y = ConcreteRNumber(x)
 
     f = Reactant.compile((y,)) do z
-        z + Reactant.TracedUtils.promote_to(Reactant.TracedRNumber{ComplexF64}, 1.0 - 3.0im)
+        z + Reactant.TracedUtils.promote_to(
+            Reactant.TracedRNumber{ComplexF32}, ComplexF32(1.0 - 3.0im)
+        )
     end
 
-    @test isapprox(f(y), 2.0 - 1.0im) broken = RunningOnTPU
+    @test isapprox(f(y), ComplexF32(2.0 - 1.0im))
 end
 
 @testset "complex reduction" begin
     x = randn(ComplexF32, 10, 10)
     x_ra = Reactant.to_rarray(x)
-    @test @jit(sum(abs2, x_ra)) ≈ sum(abs2, x) broken = RunningOnTPU
+    @test @jit(sum(abs2, x_ra)) ≈ sum(abs2, x)
 end
 
 @testset "create complex numbers" begin
     x = randn(ComplexF32)
     x_ra = Reactant.to_rarray(x; track_numbers=true)
-    @test @jit(Complex(x_ra)) == x_ra broken = RunningOnTPU
+    @test @jit(Complex(x_ra)) == x_ra
 
     x = randn(Float32)
     y = randn(Float64)
     x_ra = Reactant.to_rarray(x; track_numbers=true)
     y_ra = Reactant.to_rarray(y; track_numbers=true)
-    @test @jit(Complex(x_ra, y_ra)) == Complex(x, y) broken = RunningOnTPU
-    @test @jit(Complex(x_ra, y)) == Complex(x, y) broken = RunningOnTPU
-    @test @jit(Complex(x, y_ra)) == Complex(x, y) broken = RunningOnTPU
-    @test @jit(Complex(x_ra)) == Complex(x) == @jit(Complex(x_ra, 0)) broken = RunningOnTPU
+    @test @jit(Complex(x_ra, y_ra)) == Complex(x, y) skip = RunningOnTPU
+    @test @jit(Complex(x_ra, y)) == Complex(x, y) skip = RunningOnTPU
+    @test @jit(Complex(x, y_ra)) == Complex(x, y) skip = RunningOnTPU
+    @test @jit(Complex(x_ra)) == Complex(x) == @jit(Complex(x_ra, 0))
 end

test/optimize_comm.jl

@@ -37,6 +37,7 @@ if length(addressable_devices) ≥ 8
         @test !contains(hlo, "all-gather")
         @test contains(hlo, "collective-permute")
 
+        rotate(x)
         @jit shardy_passes = :to_mhlo_shardings rotate(rx)
         @test all(x .== convert(Array, rx))
     end

test/runtests.jl

@@ -16,10 +16,10 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         # @info "Layout tests finished"
         # @safetestset "Tracing" include("tracing.jl")
         # @info "Tracing tests finished"
-        # @safetestset "Basic" include("basic.jl") # TODO: needs fixing -- stalling currently
-        # @info "Basic tests finished"
-        @safetestset "Constructor" include("constructor.jl")
-        @info "Constructor tests finished"
+        @safetestset "Basic" include("basic.jl")
+        @info "Basic tests finished"
+        # @safetestset "Constructor" include("constructor.jl")
+        # @info "Constructor tests finished"
         @safetestset "Autodiff" include("autodiff.jl")
         @info "Autodiff tests finished"
         @safetestset "Complex" include("complex.jl")