Fix enzyme.make_zero! (#1544)

* Fix enzyme.make_zero

* fix

* with test

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* allowscalar

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: wsmoses

src/Enzyme.jl

@@ -5,6 +5,46 @@ const enzyme_dupnoneed = 3
 const enzyme_outnoneed = 4
 const enzyme_constnoneed = 5
 
+@inline function Enzyme.make_zero(x::RNumber)
+    return zero(Core.Typeof(x))
+end
+
+@inline function Enzyme.make_zero(x::RArray{FT,N})::RArray{FT,N} where {FT<:AbstractFloat,N}
+    return Base.zero(x)
+end
+
+@inline function Enzyme.make_zero(
+    x::RArray{Complex{FT},N}
+)::RArray{Complex{FT},N} where {FT<:AbstractFloat,N}
+    return Base.zero(x)
+end
+
+macro register_make_zero_inplace(sym)
+    quote
+        @inline function $sym(prev::RArray{T,N})::Nothing where {T<:AbstractFloat,N}
+            $sym(prev, nothing)
+            return nothing
+        end
+
+        @inline function $sym(prev::RArray{T,N}, seen::ST)::Nothing where {T,N,ST}
+            if Enzyme.Compiler.guaranteed_const_nongen(T, nothing)
+                return nothing
+            end
+            if !isnothing(seen)
+                if prev in seen
+                    return nothing
+                end
+                push!(seen, prev)
+            end
+            fill!(prev, zero(T))
+            return nothing
+        end
+    end
+end
+
+@register_make_zero_inplace(Enzyme.make_zero!)
+@register_make_zero_inplace(Enzyme.remake_zero!)
+
 function Enzyme.make_zero(
     ::Type{RT}, seen::IdDict, prev::RT, ::Val{copy_if_inactive}=Val(false)
 )::RT where {copy_if_inactive,RT<:Union{RArray,RNumber}}

test/autodiff.jl

@@ -231,6 +231,13 @@ vector_forward_ad(x) = Enzyme.autodiff(Forward, fn, BatchDuplicated(x, Enzyme.on
     @test res[1][4] ≈ res_enz[1][4]
 end
 
+@testset "make_zero!" begin
+    x = Reactant.to_rarray([3.1])
+    @jit Enzyme.make_zero!(x)
+
+    @test @allowscalar x[1] ≈ 0.0
+end
+
 function simple_forward(x, st)
     rng = copy(st.rng)
     y = similar(x)