Skip to content

Commit 1274e8d

Browse files
committed
Correct behavior
1 parent 0cf14f4 commit 1274e8d

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
@@ -416,8 +416,9 @@ end
416416
@test z2d_compiled.c.a === z2d_fallback.c.a
417417
@test z2d_compiled.c.b === z2d_fallback.c.b
418418

419-
# empty arrays => NoTangent()
420-
@test @inferred(Zygote.z2d(ones(1, 0), ones(16, 0))) === NoTangent()
419+
# empty dx => returns the dx
420+
@test @inferred(Zygote.z2d(ones(1, 0), ones(16, 0))) === ones(1, 0)
421+
@test @inferred(Zygote.z2d(Union{Nothing, Float64}[], ones(16, 0))) === Union{Nothing, Float64}[]
421422
end
422423

423424
@testset "ChainRules translation" begin

0 commit comments

Comments
 (0)