Use eltype to determine the proper zero to fill! in _zerolike_writeat

[SBuercklin]

Fixes, e.g, the example below:

using Unitful, UnitfulChainRules
using Zygote

Zygote.gradient(A -> ustrip.(maximum(A)), [1.0, 2.0, 3.0] * u"m")
# ERROR: DimensionError: m^-1 and 0 are not dimensionally compatible.

This occurs because we hardcode the 0 in _zerolike_writeat, which ends up getting called through the maximum rrule, among others. Using zero(eltype(dy)) we can fill with the proper 0-like.

I included a check just in case eltype(dy) == Any, in which case we default to the previous behavior as zero(Any) is undefined.

Is there a better way to determine when zero(eltype(dy)) wouldn't be valid? Maybe defaulting to ZeroTangent() would be better in the case where eltype doesn't concretely define a zero?

[mzgubic]

Thanks!

Is there a better way to determine when zero(eltype(dy)) wouldn't be valid?

I agree it is not ideal at the moment, but it is an improvement over the status quo, so I am happy for it to be merged as is. Looks like zero works for Unions of Number subtypes, and I couldn't really think of any other edge cases where it would fail. Something like

_zerofill = try zero(eltype(dy)) catch e 0 end

would be better, it doesn't harm inference, but I think Zygote will fail to do a second order derivative through this.

Maybe defaulting to ZeroTangent() would be better in the case where eltype doesn't concretely define a zero?

I'd prefer the current implementation since we likely don't want something like this:

julia> ChainRules._zerolike_writeat([1, 2, 3.0], Any[1.0], 1, 3)
3-element Vector{Any}:
  ChainRulesCore.ZeroTangent()
  ChainRulesCore.ZeroTangent()
 1.0