fix: correct type for ad annotation

Author: avik-pal

src/Enzyme.jl

@@ -295,6 +295,16 @@ function act_attr(val)
     return MLIR.IR.Attribute(val)
 end
 
+function overload_autodiff(
+    mode::CMode, f::FA, args::Vararg{Annotation,Nargs}
+) where {CMode<:Mode,FA<:Annotation,Nargs}
+    # need to guess the correct activity here. Execute the function, we will DCE it
+    res = call_with_reactant(f.val, [x.val for x in args]...)
+    return overload_autodiff(
+        mode, f, Enzyme.guess_activity(Core.Typeof(res), mode), args...
+    )
+end
+
 function overload_autodiff(
     ::CMode, f::FA, ::Type{A}, args::Vararg{Annotation,Nargs}
 ) where {CMode<:Mode,FA<:Annotation,A<:Annotation,Nargs}

src/Overlay.jl

@@ -9,6 +9,18 @@
 end
 
 # Enzyme.jl overlays
+@reactant_overlay @noinline function Enzyme.autodiff_deferred(
+    rmode::Enzyme.Mode, f::FA, args::Vararg{Annotation,Nargs}
+) where {FA<:Annotation,Nargs}
+    return overload_autodiff(rmode, f, args...)
+end
+
+@reactant_overlay @noinline function Enzyme.autodiff(
+    rmode::Enzyme.Mode, f::FA, args::Vararg{Annotation,Nargs}
+) where {FA<:Annotation,Nargs}
+    return overload_autodiff(rmode, f, args...)
+end
+
 @reactant_overlay @noinline function Enzyme.autodiff_deferred(
     rmode::Enzyme.Mode, f::FA, rt::Type{A}, args::Vararg{Annotation,Nargs}
 ) where {FA<:Annotation,A<:Annotation,Nargs}

test/autodiff.jl

@@ -2,6 +2,8 @@ using Enzyme, Reactant, Test, Random
 
 square(x) = x * 2
 
+sum_without_activity(x) = sum(abs2, x; dims=(2, 3))
+
 fwd(Mode, RT, x, y) = Enzyme.autodiff(Mode, square, RT, Duplicated(x, y))
 
 @testset "Activity" begin
@@ -39,6 +41,16 @@ fwd(Mode, RT, x, y) = Enzyme.autodiff(Mode, square, RT, Duplicated(x, y))
         Enzyme.Duplicated
 end
 
+@testset "Correct Activity Guess" begin
+    x = Reactant.to_rarray(Reactant.TestUtils.construct_test_array(Float32, 3, 4, 5, 6))
+    bx = Reactant.to_rarray(Reactant.TestUtils.construct_test_array(Float32, 3, 4, 5, 6))
+
+    res = only(@jit(Enzyme.autodiff(Forward, sum_without_activity, Duplicated(x, bx))))
+    @test res isa Reactant.ConcreteRArray{Float32,4}
+    @test size(res) == (3, 1, 1, 6)
+    @test 2 .* sum_without_activity(bx) ≈ res
+end
+
 @testset "Basic Forward Mode" begin
     res1 = @jit(
         fwd(