Prefer numeric zero over ZeroTangent for numeric arrays

FluxML · Sep 15, 2024 · 24edc21 · 24edc21
1 parent 0759953
commit 24edc21
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/compiler/chainrules.jl b/src/compiler/chainrules.jl
@@ -162,6 +162,7 @@ end
 # For arrays, whitelist the safe ones, but always look inside Any[]:
 @inline wrap_chainrules_input(dxs::AbstractArray{<:Number}) = dxs
 @inline wrap_chainrules_input(dxs::AbstractArray{<:AbstractArray{<:Number}}) = dxs
+@inline wrap_chainrules_input(dxs::AbstractArray{<:Union{Nothing,T}}) where T <: Number = map(x -> x === nothing ? zero(T) : x, dxs)
 @inline wrap_chainrules_input(dxs::AbstractArray) = map(wrap_chainrules_input, dxs)
 
 #=