Correct behavior

FluxML · Oct 7, 2024 · 8d24a29 · 8d24a29
1 parent 5c6db01
commit 8d24a29
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/compiler/chainrules.jl b/src/compiler/chainrules.jl
@@ -282,7 +282,7 @@ z2d(::Tuple{Vararg{Nothing}}, ::Tuple) = NoTangent()  # collapse all-zero case
 z2d(dx, ::Any) = dx
 z2d(dx::AbstractArray{<:Number}, primal::AbstractArray) = dx
 z2d(dx::AbstractArray{<:AbstractArray{<:Number}}, primal::AbstractArray) = dx
-z2d(dx::AbstractArray, primal::AbstractArray) = isempty(dx) ? NoTangent() : map(Zygote.z2d, dx, primal)
+z2d(dx::AbstractArray, primal::AbstractArray) = isempty(dx) ? dx : map(Zygote.z2d, dx, primal)
 
 #=
 # As an optimisation, we can convert by `reinterpret` for bitstypes, e.g. arrays of tuples of numbers

diff --git a/test/chainrules.jl b/test/chainrules.jl
@@ -420,8 +420,9 @@ end
     @test z2d_compiled.c.a === z2d_fallback.c.a
     @test z2d_compiled.c.b === z2d_fallback.c.b
 
-    # empty arrays => NoTangent()
-    @test @inferred(Zygote.z2d(ones(1, 0), ones(16, 0))) === NoTangent()
+    # empty dx => returns the dx
+    @test @inferred(Zygote.z2d(ones(1, 0), ones(16, 0))) === ones(1, 0)
+    @test @inferred(Zygote.z2d(Union{Nothing, Float64}[], ones(16, 0))) === Union{Nothing, Float64}[]
 end
 
 @testset "ChainRules translation" begin