Canonicalize the Composite before changing it to a NamedTuple

[mzgubic]

Closes #922.

The issue occurs when the following are all true:

1. We have a struct with two differentiable fields
2. We need to accumulate the gradients with respect to both fields
3. The gradients that need to be accumulated both originate from an rrule

In that case, the two gradients that originate from an rrule are Composites with disjoint fields. When they are transformed to a NamedTuple, the Zygote internal representation of derivatives w.r.t. structs, these two NamedTuples have disjoint sets of keys, and are not accumulated correctly.

By adding canonicalize(x), the Composite gets explicit Zero() fields, which means the resulting NamedTuples have the complete set of fields.

- blocked by JuliaDiff/ChainRules.jl#390 (tests only)
- need to change compat once the above is merged

[mzgubic]

Looks like we either need JuliaDiff/ChainRulesCore.jl#321 or use

    maybe_canonicalized = isabstracttype(P) ? x : canonicalize(x)
    xp = map(wrap_chainrules_output, maybe_canonicalized)

instead. It seems to be able to do type inference alright.

[DhairyaLGandhi]

Seems a bit hairy. Could you give me an idea of what canonicalize does concretely?

[mzgubic]

Yeah I'd prefer the issue but there might be reasons not to do that (I can't think of any, but there might be).

To understand canonicalize one needs to know that Composites, which are the differential types for structs, only store non-zero gradients explicitly. For example, if I have a MyStruct with two fields, a and b, then the rrule for getproperty(mystruct, a) can return Composite{MyStruct}(;a=ȳ) rather than Composite{MyStruct}(;a=ȳ, b=Zero()).

The zeros for other fields are implicit, meaning it is possible to sum Composites easily:

julia> Composite{MyStruct}(;a=ȳ) + Composite{MyStruct}(;b=ȳ2)
Composite{MyStruct}(;a=ȳ, b=ȳ2)

canonicalize will simply make these implicit Zero()s explicit, i.e.

julia> canonicalize(Composite{MyStruct}(;a=ȳ))
Composite{MyStruct}(;a=ȳ, b=Zero())

This relies on knowing the primal type (MyStruct) that Composite stores the derivative for in order to access the fields. Here, the issue is that in higher-order AD sometimes the primal type is unknown*, and in that case it is not possible to know which fields should explicitly be set to zero.

*I think we can get rid of this once Zygote uses ChainRules differential types internally

[oxinabox]

I am pretty sure we should do this.
Since otherwise the NamedTuple that comes out of a ChainRules rule will not ensc have the same fields as one from Zygote working it out on its own.
And that will lead to badness.

I commented on this here: #861 (comment)

[mzgubic]

I've added ChainRulesCore as a direct dependency to be able to specify a compat bound. I don't have merge rights but there is nothing I want to add, so please go ahead once you think it's ready

[oxinabox]

the GPU tests for o Cholesky weem flaky

See if they work for bores

Bors r+

[mzgubic]

Same test failed in a different PR which only added another test. Not sure what to think about that, but I don't think it's related to either of the PRs?

[bors]

Merge conflict.

[mzgubic]

Could you try again please?

[oxinabox]

bors r+

[bors]

Build succeeded:

• buildkite/zygote-dot-jl