Throw error when trying to accum two NamedTuples with different keys

[mzgubic]

Would have helped here:
#922 (comment)

[mzgubic]

@willtebbutt mind taking a look?

[DhairyaLGandhi]

Playing devil's advocate a bit here, but only to take the conversation forward. Should we not have it so that keys of y are a subset of keys of x?

We should document this with tests too.

[mzgubic]

Quite possibly, I don't know what the precise assumptions are here. My understanding was that Zygote internally uses NamedTuples with explicit nothing for all zero gradients? (Unlike ChainRules where Composite only needs non zero derivatives explicitly, and the others are zero implicitly)

[willtebbutt]

Playing devil's advocate a bit here, but only to take the conversation forward. Should we not have it so that keys of y are a subset of keys of x?

Ah, yeah, based on the current functionality, I think you're correct @DhairyaLGandhi .

@mzgubic I think this looks good modulo the subset stuff (I'm just assuming you'll write some tests). While you're here, would you mind renaming the argument of grad to something other than x? Having two things called x in such close proximity makes for hard reading.

[mzgubic]

Oh, when you say assumptions you just mean what this function does, rather than the broader Zygote assumptions about NamedTuples as gradients for structs? In that case I totally agree.

[willtebbutt]

LGTM. @DhairyaLGandhi should also check he's happy though.

[mzgubic]

Great, thanks, just FYI I don't have merge rights so one of you is going to have to merge it

[DhairyaLGandhi]

Looks fine, thanks!

bors r+

[bors]

Build succeeded:

• buildkite/zygote-dot-jl

[marius311]

Looks like this PR broke CI for my package CMBLensing.jl https://github.com/marius311/CMBLensing.jl/runs/2282205515

Glancing over this and linked issues, my understanding is Zygote was returning an incorrect result before, and now this errors (fwiw, previously I wasn't finding any incorrect gradients)

More importaintly, would one of you guys be able to offer some hint as to how to fix whatever my package presumably was doing wrong before? I'm not expecting anyone to go read my source code, but just wondering if you could suggest the general thing I should be looking for?

[DhairyaLGandhi]

Well it wasn't incorrect in the sense that things that shouldn't have gradients, never had gradients to begin with.

Generally you want to check for cases where you were generating namedtuples in custom adjoints.

[mzgubic]

From the error message

ArgumentError:
NamedTuple{(:op, :recompute_function, :parameters),Tuple{Diagonal{Complex{Float64},CMBLensing.BaseField{CMBLensing.Basis2Prod{CMBLensing.𝐄𝐁,Fourier},ProjLambert{Float64,Array{Float64,1},Array{Float64,2}},Complex{Float64},Array{Complex{Float64},3}}},Nothing,Nothing}}
keys must be a subset of 
NamedTuple{(:L,),Tuple{NamedTuple{(:op, :recompute_function, :parameters),Tuple{Diagonal{Complex{Float64},CMBLensing.BaseField{CMBLensing.Basis2Prod{CMBLensing.𝐄𝐁,Fourier},ProjLambert{Float64,Array{Float64,1},Array{Float64,2}},Complex{Float64},Array{Complex{Float64},3}}},Nothing,Nothing}}}}
keys

It looks like the accumulation is between a NamedTuple and a nested NamedTuple. It looks like previously the contents of the non-nested NamedTuple were ignored, which may have given the correct result if the values were zero.

[marius311]

Thanks a ton for the replies @DhairyaLGandhi @mzgubic, ended up finding the mistake, it was definitely me returning an incorrect adjoint, and your tips were super helpful in tracking it down!

[mzgubic]

Great! Is that the one?marius311/CMBLensing.jl@efb618d

[marius311]

Yep, good sleuthing :) If you're curious / confused whats going on there, it originates from somethings like this : https://discourse.julialang.org/t/how-to-deal-with-zygote-sometimes-pirating-its-own-adjoints-with-worse-ones/32579/2?u=marius311

[mzgubic]

Oh, I thought it was just that logdet has two arguments and needs a pullback of a function with two arguments (rather than a one argument closure as before). Well, it works now, so we don't have to dig in deeper :)

[marius311]

No, you're exactly right as to what the mistake was c.f. this PR, although the underlying reason there's a lot of explicit Zygote.pullback calls in that file has to do with that forum post.