Support ChainRulesCore.Tangent?

[oschulz]

I realized that currently, Functors doesn't support ChainRulesCore.Tangent:

using Functors, ChainRulesCore
x = (a = 2, b = 3)
dx = Tangent{typeof(x)](a = 4, b = 9)
Functors.functor(dx)

results in

((), Functors.var"#1#2"{Tangent{NamedTuple{(:a, :b), Tuple{Int64, Int64}}, NamedTuple{(:a, :b), Tuple{Int64, Int64}}}}(Tangent{NamedTuple{(:a, :b), Tuple{Int64, Int64}}}(a = 4, b = 9)))

Typical use cases of Functors will usually use re of x instead of dx of course - still, it might be nice to support functor for tangents.

[darsnack]

This is something ChainRulesCore.jl has to do by calling @functor Tangent. Functors.jl itself doesn't define any functors.

cc @oxinabox

[DhairyaLGandhi]

This is where the NamedTuple versions make a lot more sense since those are implicitly functors themselves. Or at least have the same underlying tree as the original.

[ToucheSir]

This is the dual of FluxML/Zygote.jl#1042. One could imagine a Functor{T} type that works much like Tangent, but the question is whether the two are similar enough that they can be unified. Either way, the type will require a custom functor implementation because the list of fields is dynamic.

[oschulz]

I guess I should move this issue to ChainRulesCore - I thought it might be doable with some generic magic, like Functors does for some other things - if it needs a custom functor implementation, that should go into CRC, or course.

[oschulz]

Moved to JuliaDiff/ChainRulesCore.jl#439

[DhairyaLGandhi]

Well, since Tangent is a very generic container of sorts, one can imagine it does share the same representation of the tree as regular models would. That is, after all, its job. It may only have incomplete information because not all fields would have a corresponding bit in the Tangent, but that can be handled in the function doing the operation itself. At the very least, one would expect that all the leaves are present.

We make use of the fact that the trees are similar to do structural optimisation in the upcoming DDP stuff too and that works with the way zygote returns the gradients just fine, so a tangent also ought to be able to.

[oschulz]

Just curious, what upcoming DDP stuff?

[ToucheSir]

https://github.com/DhairyaLGandhi/DaggerFlux.jl

[oxinabox]

Well, since Tangent is a very generic container of sorts, one can imagine it does share the same representation of the tree as regular models would. That is, after all, it's job. It may only have incomplete information because not all fields would have a corresponding bit in the Tangent

If you call canonicalize on a Tangent then it will fill in all the fields that don't have a tangent in them with ZeroTangent().
So it will have same structure.