test: remove randomness from testing (#1808)

* test: remove randomness from testing

* test: more randomness + remove enzymellvm from testing

* test: more

* test: more

* test: remove enzyme testing with finite diff

* test: revert linalg

Author: avik-pal

src/Reactant.jl

@@ -351,6 +351,9 @@ function set_default_backend(backend::Union{String,XLA.AbstractClient})
     return nothing
 end
 
+# Not part of the public API. Exclusively for testing purposes.
+include("TestUtils.jl")
+
 include("Precompile.jl")
 
 end # module

src/TestUtils.jl

@@ -0,0 +1,41 @@
+module TestUtils
+
+using ..Reactant: Reactant, TracedRArray
+using ReactantCore: ReactantCore
+using LinearAlgebra: LinearAlgebra
+
+function construct_test_array(::Type{T}, dims::Int...) where {T<:AbstractFloat}
+    flat_vector = collect(T, 1:prod(dims))
+    flat_vector ./= prod(dims)
+    return reshape(flat_vector, dims...)
+end
+
+function construct_test_array(::Type{Complex{T}}, dims::Int...) where {T<:AbstractFloat}
+    flat_vector = collect(T, 1:prod(dims))
+    flat_vector ./= prod(dims)
+    return reshape(complex.(flat_vector, flat_vector), dims...)
+end
+
+function construct_test_array(::Type{T}, dims::Int...) where {T}
+    return reshape(collect(T, 1:prod(dims)), dims...)
+end
+
+function finite_difference_gradient(
+    f, x::AbstractArray{T}; epsilon=eps(T)^(3 / 4)
+) where {T}
+    onehot_matrix = Reactant.promote_to(
+        TracedRArray{Reactant.unwrapped_eltype(T),2}, LinearAlgebra.I(length(x))
+    )
+    perturbation = reshape(onehot_matrix .* epsilon, size(x)..., length(x))
+    f_input = cat(x .+ perturbation, x .- perturbation; dims=ndims(x) + 1)
+
+    f_evaluated = mapslices(f, f_input; dims=ntuple(identity, ndims(x)))
+    return ReactantCore.materialize_traced_array(
+        reshape(
+            (f_evaluated[1:length(x)] - f_evaluated[(length(x) + 1):end]) ./ (2 * epsilon),
+            size(x),
+        ),
+    )
+end
+
+end

test/autodiff.jl

@@ -40,10 +40,6 @@ fwd(Mode, RT, x, y) = Enzyme.autodiff(Mode, square, RT, Duplicated(x, y))
 end
 
 @testset "Basic Forward Mode" begin
-    ores1 = fwd(Forward, Duplicated, ones(3, 2), 3.1 * ones(3, 2))
-    @test typeof(ores1) == NamedTuple{(Symbol("1"),),Tuple{Array{Float64,2}}}
-    @test ores1[1] ≈ 6.2 * ones(3, 2)
-
     res1 = @jit(
         fwd(
             Forward,
@@ -54,11 +50,7 @@ end
     )
 
     @test res1 isa Tuple{<:ConcreteRArray{Float64,2}}
-    @test res1[1] ≈ ores1[1]
-
-    ores1 = fwd(ForwardWithPrimal, Duplicated, ones(3, 2), 3.1 * ones(3, 2))
-    @test typeof(ores1) ==
-        NamedTuple{(Symbol("1"), Symbol("2")),Tuple{Array{Float64,2},Array{Float64,2}}}
+    @test res1[1] ≈ fill(6.2, 3, 2)
 
     res1 = @jit(
         fwd(
@@ -70,11 +62,8 @@ end
     )
 
     @test res1 isa Tuple{<:ConcreteRArray{Float64,2},<:ConcreteRArray{Float64,2}}
-    @test res1[1] ≈ ores1[1]
-    @test res1[2] ≈ ores1[2]
-
-    ores1 = fwd(Forward, Const, ones(3, 2), 3.1 * ones(3, 2))
-    @test typeof(ores1) == Tuple{}
+    @test res1[1] ≈ fill(6.2, 3, 2)
+    @test res1[2] ≈ fill(2.0, 3, 2)
 
     res1 = @jit(
         fwd(
@@ -87,9 +76,6 @@ end
 
     @test typeof(res1) == Tuple{}
 
-    ores1 = fwd(ForwardWithPrimal, Const, ones(3, 2), 3.1 * ones(3, 2))
-    @test typeof(ores1) == NamedTuple{(Symbol("1"),),Tuple{Array{Float64,2}}}
-
     res1 = @jit(
         fwd(
             set_abi(ForwardWithPrimal, Reactant.ReactantABI),
@@ -100,7 +86,7 @@ end
     )
 
     @test res1 isa Tuple{<:ConcreteRArray{Float64,2}}
-    @test res1[1] ≈ ores1[1]
+    @test res1[1] ≈ fill(2.0, 3, 2)
 end
 
 function gw(z)
@@ -140,7 +126,7 @@ function cached_return(x, stret::StateReturn1)
 end
 
 @testset "Cached Return: Issue #416" begin
-    x = rand(10)
+    x = Reactant.TestUtils.construct_test_array(Float64, 10)
     x_ra = Reactant.to_rarray(x)
 
     stret = StateReturn(nothing)
@@ -187,7 +173,7 @@ end
 end
 
 @testset "onehot" begin
-    x = Reactant.to_rarray(rand(3, 4))
+    x = Reactant.to_rarray(ones(3, 4))
     hlo = @code_hlo optimize = false Enzyme.onehot(x)
     @test !contains("stablehlo.constant", repr(hlo))
 end
@@ -202,27 +188,26 @@ end
     x = reshape(collect(Float32, 1:6), 3, 2)
     x_ra = Reactant.to_rarray(x)
     res = @jit vector_forward_ad(x_ra)
-    res_enz = vector_forward_ad(x)
 
     @test x_ra ≈ x # See https://github.com/EnzymeAD/Reactant.jl/issues/1733
-    @test res[1][1] ≈ res_enz[1][1]
-    @test res[1][2] ≈ res_enz[1][2]
-    @test res[1][3] ≈ res_enz[1][3]
-    @test res[1][4] ≈ res_enz[1][4]
-    @test res[1][5] ≈ res_enz[1][5]
-    @test res[1][6] ≈ res_enz[1][6]
+    @test res[1][1] ≈ 2
+    @test res[1][2] ≈ 4
+    @test res[1][3] ≈ 6
+    @test res[1][4] ≈ 8
+    @test res[1][5] ≈ 10
+    @test res[1][6] ≈ 12
 
     oh = Enzyme.onehot(x)
     oh_stacked = stack(oh)
     oh_ra = Reactant.to_rarray(oh_stacked)
     res2 = @jit vector_forward_ad2(x_ra, oh_ra)
 
-    @test res2[1][1] ≈ res_enz[1][1]
-    @test res2[1][2] ≈ res_enz[1][2]
-    @test res2[1][3] ≈ res_enz[1][3]
-    @test res2[1][4] ≈ res_enz[1][4]
-    @test res2[1][5] ≈ res_enz[1][5]
-    @test res2[1][6] ≈ res_enz[1][6]
+    @test res2[1][1] ≈ 2
+    @test res2[1][2] ≈ 4
+    @test res2[1][3] ≈ 6
+    @test res2[1][4] ≈ 8
+    @test res2[1][5] ≈ 10
+    @test res2[1][6] ≈ 12
 end
 
 function fn2!(y, x)
@@ -245,34 +230,27 @@ end
     dx3_ra = Reactant.to_rarray(dx3)
     dx4_ra = Reactant.to_rarray(dx4)
 
-    dy1 = zeros(2)
-    dy2 = zeros(2)
-    dy3 = zeros(2)
-    dy4 = zeros(2)
+    dy1 = ones(2) .* 1
+    dy2 = ones(2) .* 2
+    dy3 = ones(2) .* 3
+    dy4 = ones(2) .* 4
     dy1_ra = Reactant.to_rarray(dy1)
     dy2_ra = Reactant.to_rarray(dy2)
     dy3_ra = Reactant.to_rarray(dy3)
     dy4_ra = Reactant.to_rarray(dy4)
 
-    autodiff(
-        ReverseWithPrimal,
-        fn2!,
-        BatchDuplicated(y, (dy1, dy2, dy3, dy4)),
-        BatchDuplicated(x, (dx1, dx2, dx3, dx4)),
-    )
-
     @jit autodiff(
         Reverse,
         fn2!,
         BatchDuplicated(y_ra, (dy1_ra, dy2_ra, dy3_ra, dy4_ra)),
         BatchDuplicated(x_ra, (dx1_ra, dx2_ra, dx3_ra, dx4_ra)),
     )
 
-    @test y ≈ y_ra
-    @test dy1 ≈ dy1_ra
-    @test dy2 ≈ dy2_ra
-    @test dy3 ≈ dy3_ra
-    @test dy4 ≈ dy4_ra
+    @test y_ra ≈ x .^ 2
+    @test dx1_ra ≈ 2 .* x .* dy1
+    @test dx2_ra ≈ 2 .* x .* dy2
+    @test dx3_ra ≈ 2 .* x .* dy3
+    @test dx4_ra ≈ 2 .* x .* dy4
 end
 
 @testset "make_zero!" begin
@@ -300,7 +278,7 @@ function gradient_fn(x, st)
 end
 
 @testset "seed" begin
-    x = Reactant.to_rarray(rand(2, 2))
+    x = Reactant.to_rarray(Reactant.TestUtils.construct_test_array(Float64, 2, 2))
     st = (; rng=Reactant.ReactantRNG())
 
     @test begin
@@ -344,7 +322,7 @@ function zero_grad2(x)
 end
 
 @testset "ignore_derivatives" begin
-    x = Reactant.to_rarray(rand(Float32, 4, 4))
+    x = Reactant.to_rarray(Reactant.TestUtils.construct_test_array(Float32, 4, 4))
 
     res1 = @jit Enzyme.gradient(Reverse, simple_grad_without_ignore, x)
     @test res1[1] ≈ (2 .* Array(x) .+ 4)