diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl index e4ef4c1f..5dc9032f 100644 --- a/ext/RecursiveArrayToolsZygoteExt.jl +++ b/ext/RecursiveArrayToolsZygoteExt.jl @@ -4,6 +4,7 @@ using RecursiveArrayTools using Zygote using Zygote: FillArrays, ChainRulesCore, literal_getproperty, @adjoint +using Zygote.ChainRulesCore: AbstractZero function ChainRulesCore.rrule( T::Type{<:RecursiveArrayTools.GPUArraysCore.AbstractGPUArray}, @@ -90,7 +91,7 @@ end function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractVectorOfArray) m = map(enumerate(d)) do (idx, d_i) - isnothing(d_i) && return zero(A.u[idx]) + (isnothing(d_i) || d_i isa AbstractZero) && return zero(A.u[idx]) d_i end return VectorOfArray(m) @@ -98,7 +99,7 @@ end function vofa_u_adjoint(d, A::RecursiveArrayTools.AbstractDiffEqArray) m = map(enumerate(d)) do (idx, d_i) - isnothing(d_i) && return zero(A.u[idx]) + (isnothing(d_i) || d_i isa AbstractZero) && return zero(A.u[idx]) d_i end return DiffEqArray(m, A.t) @@ -106,7 +107,7 @@ end @adjoint function literal_getproperty(A::ArrayPartition, ::Val{:x}) function literal_ArrayPartition_x_adjoint(d) - (ArrayPartition((isnothing(d[i]) ? zero(A.x[i]) : d[i] for i in 1:length(d))...),) + (ArrayPartition((isnothing(d[i]) || d[i] isa AbstractZero ? zero(A.x[i]) : d[i] for i in 1:length(d))...),) end A.x, literal_ArrayPartition_x_adjoint end diff --git a/test/adjoints.jl b/test/adjoints.jl index 4e60c108..4bb1a644 100644 --- a/test/adjoints.jl +++ b/test/adjoints.jl @@ -1,4 +1,5 @@ using RecursiveArrayTools, Zygote, ForwardDiff, ReverseDiff, Test +using Zygote: ChainRulesCore function loss(x) return sum(abs2, Array(VectorOfArray([x .* i for i in 1:5]))) @@ -99,3 +100,21 @@ end let g = ReverseDiff.gradient(x -> sum(VectorOfArray([VectorOfArray([x])])), float.(1:3)) @test g ≈ ones(3) end + +# A `ZeroTangent`/`AbstractZero` slice in the cotangent (structural zero gradient) +# must be treated like an absent gradient instead of erroring on `size(::ZeroTangent)`. +let ext = Base.get_extension(RecursiveArrayTools, :RecursiveArrayToolsZygoteExt) + voa = VectorOfArray([Float64[i, i, i] for i in 1:3]) + dea = DiffEqArray([Float64[i, i, i] for i in 1:3], 1:3) + d = Any[Float64[1, 1, 1], ChainRulesCore.ZeroTangent(), Float64[3, 3, 3]] + expected = [Float64[1, 1, 1], Float64[0, 0, 0], Float64[3, 3, 3]] + + g_voa = ext.vofa_u_adjoint(d, voa) + @test g_voa isa VectorOfArray + @test g_voa.u == expected + + g_dea = ext.vofa_u_adjoint(d, dea) + @test g_dea isa DiffEqArray + @test g_dea.u == expected + @test g_dea.t == dea.t +end