fix: finite diff gradient accidental type promotion

Author: avik-pal

src/TestUtils.jl

@@ -21,12 +21,13 @@ function construct_test_array(::Type{T}, dims::Int...) where {T}
 end
 
 function finite_difference_gradient(
-    f, x::AbstractArray{T}; epsilon=eps(T)^(3 / 4)
+    f, x::AbstractArray{T}; epsilon=eps(T)^(T(3 / 4))
 ) where {T}
     onehot_matrix = Reactant.promote_to(
-        TracedRArray{Reactant.unwrapped_eltype(T),2}, LinearAlgebra.I(length(x))
+        TracedRArray{Reactant.unwrapped_eltype(T),2},
+        LinearAlgebra.Diagonal(fill(epsilon, length(x))),
     )
-    perturbation = reshape(onehot_matrix .* epsilon, size(x)..., length(x))
+    perturbation = reshape(onehot_matrix, 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)))

test/autodiff.jl

@@ -366,3 +366,9 @@ end
 
     @test @jit(jvp_vjp_cubic(v_r, x_r, lambdas_r)) ≈ fill(6, (3, 2))
 end
+
+@testset "Finite Difference Gradient" begin
+    x = Reactant.to_rarray(Reactant.TestUtils.construct_test_array(Float16, 2, 2))
+    res = @jit Reactant.TestUtils.finite_difference_gradient(sum, x)
+    @test res isa Reactant.ConcreteRArray{Float16,2}
+end