Skip to content

Commit 8d24a29

Browse files
committed
Correct behavior
1 parent 5c6db01 commit 8d24a29

File tree

2 files changed

+4
-3
lines changed

2 files changed

+4
-3
lines changed

src/compiler/chainrules.jl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,7 @@ z2d(::Tuple{Vararg{Nothing}}, ::Tuple) = NoTangent() # collapse all-zero case
282282
z2d(dx, ::Any) = dx
283283
z2d(dx::AbstractArray{<:Number}, primal::AbstractArray) = dx
284284
z2d(dx::AbstractArray{<:AbstractArray{<:Number}}, primal::AbstractArray) = dx
285-
z2d(dx::AbstractArray, primal::AbstractArray) = isempty(dx) ? NoTangent() : map(Zygote.z2d, dx, primal)
285+
z2d(dx::AbstractArray, primal::AbstractArray) = isempty(dx) ? dx : map(Zygote.z2d, dx, primal)
286286

287287
#=
288288
# As an optimisation, we can convert by `reinterpret` for bitstypes, e.g. arrays of tuples of numbers

test/chainrules.jl

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -420,8 +420,9 @@ end
420420
@test z2d_compiled.c.a === z2d_fallback.c.a
421421
@test z2d_compiled.c.b === z2d_fallback.c.b
422422

423-
# empty arrays => NoTangent()
424-
@test @inferred(Zygote.z2d(ones(1, 0), ones(16, 0))) === NoTangent()
423+
# empty dx => returns the dx
424+
@test @inferred(Zygote.z2d(ones(1, 0), ones(16, 0))) === ones(1, 0)
425+
@test @inferred(Zygote.z2d(Union{Nothing, Float64}[], ones(16, 0))) === Union{Nothing, Float64}[]
425426
end
426427

427428
@testset "ChainRules translation" begin

0 commit comments

Comments
 (0)