accum NamedTuple sometimes silently drops derivatives

[willtebbutt]

This works correctly:

julia> Zygote.accum((a=nothing, b=4), (a=4,))
(a = 4, b = 4)

This drops the derivative w.r.t. b:

julia> Zygote.accum((a=4, ), (a=nothing, b=4))
(a = 4,)

Now this is definitely surprising, and I think it's a bug. The only situation in which it's intended behaviour is if it's the case that the derivative w.r.t. some struct should always be a NamedTuple containing all of the associated fields, since this would preclude this kind of problem.

Note that ChainRule types don't suffer from this issue:

julia> Composite{Any}(; a=4) + Composite{Any}(; a=Zero(), b=4)
Composite{Any}(a = 4, b = 4)

so we can probably nick the ChainRulesCore implementation of + for Composites and deploy it for accum with NamedTuples.

[willtebbutt]

Although looking at what Zygote does by default, perhaps it is assumed that all fields will always be present and I just missed the memo?

julia> Zygote.gradient(x -> sin(x.a), (a=4,b=4))
((a = -0.6536436208636119, b = nothing),)

(Had the gradient not had a field b, it would clearly be the case that Zygote doesn't make this assumption).

[oxinabox]

The only situation in which it's intended behaviour is if it's the case that the derivative w.r.t. some struct should always be a NamedTuple containing all of the associated fields, since this would preclude this kind of problem.

This is the case. As I understand it.
Zygote requires all structural derivatives to be NamedTuples with the same fields (in the same order?) as the primal type.
What ChainRules calls the canonical form.

As such Zygote.accum((a=nothing, b=4), (a=4,))
and Zygote.accum((a=4, ), (a=nothing, b=4))
should both error.
Since it means something went wrong somewhere to give you a bad dertivative.
Probably an incorrect rule.

Actually on that basis:
the Zygote-ChainRiles connecting code is i think missing a call to canonicalize
here

Zygote.jl/src/compiler/chainrules.jl

Lines 50 to 53 in 4fdb691

	@eval @inline function wrap_chainrules_output(x::ChainRules.Composite{P, T}) where {P, T<:$T_outer}
	xp = map(wrap_chainrules_output, x)
	convert($T_outer, xp)
	end

We should fix both this not erroring and ChainRules not getting canonicalized

[willtebbutt]

Since it means something went wrong somewhere to give you a bad dertivative.
Probably an incorrect rule.

This is exactly what happened. I returned a subset of the required fields from a particular hand-writen _pullback. Once I added in the required fields I got the right answer.

[mzgubic]

Closed by #924