fix: ignore_derivatives need to preserve structure (#1379)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.128"
+version = "0.2.129"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Enzyme.jl

@@ -549,12 +549,15 @@ be applied on a nested structure of arrays and we will apply the operation on ea
 leaves.
 """
 function ignore_derivatives(args...)
-    return map(args) do arg
-        return Functors.fmap(arg) do argᵢ
-            if argᵢ isa TracedType || argᵢ isa AnyTracedRArray
-                return Ops.ignore_derivatives(materialize_traced_array(argᵢ))
-            end
-            return argᵢ
-        end
+    res = map(ignore_derivatives_internal, args)
+    length(args) == 1 && return only(res)
+    return res
+end
+
+function ignore_derivatives_internal(arg)
+    return Functors.fmap(arg) do argᵢ
+        argᵢ isa AnyTracedRArray && (argᵢ = materialize_traced_array(argᵢ))
+        argᵢ isa TracedType && return Ops.ignore_derivatives(argᵢ)
+        return argᵢ
     end
 end

test/autodiff.jl

@@ -282,6 +282,14 @@ function simple_grad_with_ignore(x::AbstractArray{T}) where {T}
     return Reactant.ignore_derivatives(sum(x; dims=1), x .- 1, (x, x .+ 2)), sum(abs2, x)
 end
 
+function zero_grad(x)
+    return Reactant.ignore_derivatives(sum(x))
+end
+
+function zero_grad2(x)
+    return Reactant.ignore_derivatives(sum(x), x)
+end
+
 @testset "ignore_derivatives" begin
     x = Reactant.to_rarray(rand(Float32, 4, 4))
 
@@ -290,4 +298,12 @@ end
 
     res2 = @jit Enzyme.gradient(Reverse, simple_grad_with_ignore, x)
     @test res2[1] ≈ (2 .* Array(x))
+
+    ∂x, result = @jit Enzyme.gradient(ReverseWithPrimal, zero_grad, x)
+    @test result isa ConcreteRNumber{Float32}
+    @test ∂x[1] ≈ zeros(Float32, 4, 4)
+
+    ∂x2, result2 = @jit Enzyme.gradient(ReverseWithPrimal, zero_grad2, x)
+    @test result2 isa Tuple{<:ConcreteRNumber{Float32},<:ConcreteRArray{Float32,2}}
+    @test ∂x2[1] ≈ zeros(Float32, 4, 4)
 end