Add `ProjectTo(::NamedTuple)` #515

devmotion · 2021-11-10T22:58:32Z

This PR adds ProjectTo(::NamedTuple) according to the suggestion by @mcabbott (I added you as a co-author to give credit). I only added two additional more descriptive error messages and some tests, similar to the implementation for Tuples.

Fixes #511.

Co-authored-by: Michael Abbott <32575566+mcabbott@users.noreply.github.com>

devmotion · 2021-11-10T23:16:31Z

StatsFuns and Diffractor test errors seem unrelated.

mcabbott · 2021-11-10T23:57:34Z

Zygote seems to always make the whole tuple, and in the right order:

julia> gradient(x -> x.a /x.b , (a=1, b=2, c=3))
((a = 0.5, b = -0.25, c = nothing),)

But Diffractor seems to just keep the nonzero ones:

julia> g = gradient(x -> x.a /x.b , (a=1, b=2, c=3))
(Tangent{NamedTuple{(:a, :b, :c), Tuple{Int64, Int64, Int64}}}(b = -0.25, a = 0.5),)

julia> g[1].c
ChainRulesCore.ZeroTangent()

julia> ChainRulesCore.canonicalize(g[1])
Tangent{NamedTuple{(:a, :b, :c), Tuple{Int64, Int64, Int64}}}(a = 0.5, b = -0.25, c = ZeroTangent())

I suspect this means it can't just call backing and map, but will have to at least accept subsets of the keys. Whether it should also produce them or not I don't know.

mzgubic · 2021-11-11T10:27:12Z

Yeah, Zygote didn't use to do that and that led to issues, see e.g.
FluxML/Zygote.jl#922 (comment)
which was fixed by FluxML/Zygote.jl#926

Zygote has to do that because it uses NamedTuples, not Tangents (where the elements that are not there are ZeroTangent automatically). Diffractor uses Tangents internally so it should be fine.

I guess we just have to do backing(canonicalize(tangent)). ~~I don't see why we can't use map? That's a NamedTuple projected onto a NamedTuple, nothing to do with Tangent right?~~

devmotion · 2021-11-11T10:32:34Z

map requires that the names are identical. It is quite straightforward to write a map alternative that would just project the existing derivatives. Would this be preferred over using canonicalize? And is it OK if the order of the keys is different in the resulting Tangent than in the projector, or should they be reordered?

mzgubic · 2021-11-11T11:25:38Z

The order does not matter as far as I know.

I think what we want is to allow

julia> nt = (;a=1.0, b=2.0);

julia> project = ProjectTo(nt);

julia> t = Tangent{typeof(nt)}(;a=2.3)

julia> project(t)

but not

julia> project((;a=1.0))

The reason being that it is a feature of Tangents that they are implicitly ZeroTangent for all non explicitly specified fields.
On the other hand, ChainRules doesn't think a NamedTuple is a valid tangent type, so we should probably throw up a fuss if there is anything remotely suspicious going on?

Do we ever see a NamedTuple being directly projected? (As opposed to from it being the backing of a Tangent?). Should we even allow it at all?

devmotion · 2021-11-11T11:41:46Z

Do we ever see a NamedTuple being directly projected? (As opposed to from it being the backing of a Tangent?). Should we even allow it at all?

This is the case that errors in the CRTestUtils PR: FiniteDifferences returns a NamedTuple and we have to project it to the Tangent representation. Currently this is done with a custom method but it only covers Tuple and NamedTuple but eg not nested Tuples and NamedTuples. @oxinabox suggested to use ProjectTo instead of adding more dispatches to the custom method.

mcabbott · 2021-11-11T14:34:20Z

Allowing project((;a=1.0)) seems like a good idea to me. IIRC one reason the equivalent Tuple rule does accepts tuples is that it makes writing rules easier: you don't have to construct the Tangent by hand, and can treat Union{Tuple, Vector} the same way.

Perhaps it should just iterate over the keys of this input? If any of them are not found in the projector, that's an error.

devmotion · 2021-11-11T14:36:37Z

OK, I'll update the PR this evening (have a draft on my computer but are busy ATM) 🙂

codecov-commenter · 2021-11-11T16:25:02Z

Codecov Report

Merging #515 (d1bb3b2) into main (f9304b3) will increase coverage by 0.17%.
The diff coverage is 95.45%.

@@            Coverage Diff             @@
##             main     #515      +/-   ##
==========================================
+ Coverage   92.91%   93.09%   +0.17%     
==========================================
  Files          15       15              
  Lines         819      854      +35     
==========================================
+ Hits          761      795      +34     
- Misses         58       59       +1

Impacted Files	Coverage Δ
src/projection.jl	`97.29% <95.45%> (-0.11%)`	⬇️
src/rule_definition_tools.jl	`96.27% <0.00%> (+0.02%)`	⬆️
src/accumulation.jl	`97.22% <0.00%> (+0.07%)`	⬆️
src/tangent_types/thunks.jl	`95.00% <0.00%> (+0.10%)`	⬆️
src/tangent_types/tangent.jl	`85.50% <0.00%> (+0.32%)`	⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update f9304b3...d1bb3b2. Read the comment docs.

oxinabox · 2021-11-11T16:33:14Z

src/projection.jl

+    return project_type(project)(; dy...)
+end
+
+# Diffractor returns not necessarily a named tuple with all keys and of the same order as


what does Diffractor have to do with anything, and why does it return a namedtuple?
It should be a Tangent.

It refers to #515 (comment). The Tangents are already unpacked at this stage.

oxinabox · 2021-11-11T16:38:08Z

src/projection.jl

+# Diffractor returns not necessarily a named tuple with all keys and of the same order as
+# the projector
+# Thus we can't use `map`
+function _project_namedtuple(f::NamedTuple{fn,ft}, x::NamedTuple{xn,xt}) where {fn,ft,xn,xt}


what is this function for?

Can't we just stick the thing into a Tangent{typeof(f), typeof(x)}(x) ?
which should robustly handly non-present keys and keys in different orders.
And if for some reason we can't handle that then add a canonicalize ?

It is our custom projection map. Initially, in the first commit I just used map with the named tuple of projectors and named tuple of derivatives, as suggested by @mcabbott. However, map requires that the names of both named tuples are exactly identical, i.e., all derivatives are present and in the same order as the projectors. This function here just maps all existing derivatives and throws a more descriptive error if a derivative is present without corresponding projector.

I guess it could take short namedtuples and route them through Tangent -> canonise -> backing -> map -> Tangent, to re-use more stuff:

julia> using ChainRulesCore julia> x = (a=1, b=2, c=3); dx = (b=400,); julia> Tangent{typeof(x)}(; dx...) Tangent{NamedTuple{(:a, :b, :c), Tuple{Int64, Int64, Int64}}}(b = 400,) julia> ChainRulesCore.canonicalize(ans) Tangent{NamedTuple{(:a, :b, :c), Tuple{Int64, Int64, Int64}}}(a = ZeroTangent(), b = 400, c = ZeroTangent()) julia> ChainRulesCore.backing(ans) (a = ZeroTangent(), b = 400, c = ZeroTangent())

My slight reservation about all approaches really is whether we can insert enough complication to confuse Diffractor when it wants to take a 3rd derivative or something.

why do we need the backing step?
why not
Tangent -> map, which already returns a Tangent?

Well, like this it doesn't:

julia> tang Tangent{NamedTuple{(:a, :b, :c), Tuple{Int64, Int64, Int64}}}(a = ZeroTangent(), b = 400, c = ZeroTangent()) julia> projs = map(ProjectTo, x); julia> map((f,x) -> f(x), projs, tang) 3-element Vector{Any}: ZeroTangent() 400.0 ZeroTangent()

I haven't dug through all the functions closely, recently, but my reservation here is that this seems close to being a second use of canonicalize, just with a different carefully optimised generated implementation. It seems that if ever something breaks one, we'll have to fix both.

Is there a precedent anywhere else here about whether filling in all fields with NoTangent is preferable / not compared to leaving omitted ones out?

We would preferably not fill in all fields.

This was my understanding as well - and hence I don't think one should use canonicalize here since we don't want to fill all fields.

@mcabbott Are you OK with merging the PR as is and improving the implementation later, if e.g. there is a clear need for a two argument map?

Yes, do it. I think this is the right behaviour. I do wish it could be shorter but that's not the end of the world. Sorry about dragging this out so long.

devmotion · 2021-11-11T21:03:42Z

As an explanation for the current state of this PR and how ProjectTo(::NamedTuple) acts:

ProjectTo(::NamedTuple) is an analogue of of ProjectTo(::Tuple), just with Tangent{<:NamedTuple} instead of Tangent{<:Tuple} and a named tuple of projectors instead of a tuple of projectors
ProjectTo{<:Tangent{<:NamedTuple}} allows to project Tangent{<:NamedTuple} and NamedTuple to the desired space
- the latter is consistent with ProjectTo{<:Tangent{<:Tuple}} and e.g. needed in Fix tangent type of arrays of (named) tuples from FD ChainRulesTestUtils.jl#224 (I assume it is generally useful if one wants to project a named tuple to the correct return value in a custom rrule)
- to avoid code duplication the former falls back to the latter (via backing)
Since it is allowed (it seems?) to use Tangent types in rrules and from AD backends (such as Diffractor, see the example above) that do
- only contain derivatives of a subset of elements/entries in the primal input (rest assumed to be ZeroTangent)
- contain derivatives of elements/entries in the primal input in any order
we do not use map((f, x) -> f(x), backing(projector), dx) in the intermediate projection step (see @mcabbott's initial suggestion) but instead just project every entry in dx (map(f, ::NamedTuple, ::NamedTuple) requires that the names of both named tuples are identical, i.e., no subsets and only in the same order)

The projection step is implemented with a function that contains an if @generated block to ensure that the output can be inferred and it can be computed efficiently. It will display a hopefully helpful error message if any keys are present in dx for which no projector is defined.

A code example:

julia> using ChainRulesCore

julia> x = (a=1.0, b=[1.0, 3.0, 4.0], c=false);

julia> pt1 = ProjectTo(x)
ProjectTo{Tangent{NamedTuple{(:a, :b, :c),Tuple{Float64,Array{Float64,1},Bool}},T} where T}(a = ProjectTo{Float64}(), b = ProjectTo{AbstractArray}(element = ProjectTo{Float64}(), axes = (Base.OneTo(3),)), c = ProjectTo{NoTangent}())

# subsets are OK
julia> pt1(Tangent{typeof(x)}(; b = [3.0 + 0*im, 2.0, 1.0]))
Tangent{NamedTuple{(:a, :b, :c),Tuple{Float64,Array{Float64,1},Bool}}}(b = [3.0, 2.0, 1.0],)

julia> pt1((b = [3.0 + 0*im, 2.0, 1.0],))
Tangent{NamedTuple{(:a, :b, :c),Tuple{Float64,Array{Float64,1},Bool}}}(b = [3.0, 2.0, 1.0],)

# order does not matter
julia> pt1(Tangent{typeof(x)}(; b = [3.0 + 0*im, 2.0, 1.0], a=big(π)))
Tangent{NamedTuple{(:a, :b, :c),Tuple{Float64,Array{Float64,1},Bool}}}(b = [3.0, 2.0, 1.0], a = 3.141592653589793)

julia> pt1((b = [3.0 + 0*im, 2.0, 1.0], a=big(π)))
Tangent{NamedTuple{(:a, :b, :c),Tuple{Float64,Array{Float64,1},Bool}}}(b = [3.0, 2.0, 1.0], a = 3.141592653589793)

# error if derivative without projecto
julia> pt1(Tangent{typeof(x)}(; d=42, b = [3.0 + 0*im, 2.0, 1.0]))
ERROR: LoadError: ArgumentError: named tuple with keys(x) == (:a, :b, :c) cannot have a gradient with key d
...

julia> pt1((d=42, b = [3.0 + 0*im, 2.0, 1.0]))
ERROR: LoadError: ArgumentError: named tuple with keys(x) == (:a, :b, :c) cannot have a gradient with key d
...

One could "complete" the Tangent or named tuple with canonicalize, as @mcabbott suggested, since it would allow to project the derivatives more easily (no "missing" derivatives and of the same order as in the primal input). I don't know anything about the design choices behind ProjectTo but to me it seems a bit surprising to use canonicalize here if it is just fine and supported to use Tangents with a subset of keys and different order currently in rrule definitions. As long as the canonical version is not always enforced it seems a bit inconsistent to use it in ProjectTo.

Also I think the design rationale and the implementation of this PR are reasonably simple and hence I don't think canonicalize would reduce code complexity much: _project_namedtuple(projectors::NamedTuple, derivatives::NamedTuple) just projects every derivative in derivatives with the corresponding projector in projectors, or displays an error message if not existent.

src/projection.jl

devmotion · 2021-11-27T17:01:39Z

What's the status here, should the design be changed that I tried to explain above? The PR still blocks the PR in ChainRulesTestUtils.

oxinabox

Sorry, I was indecisive.
I have been wondering about https://github.com/JuliaDiff/ChainRulesCore.jl/pull/515/files#r759728574

but regardless, if we do that or not we should do this.
This PR is ok, we can always improve it later.

Add ProjectTo(::NamedTuple)

ceb1337

Co-authored-by: Michael Abbott <32575566+mcabbott@users.noreply.github.com>

Allow different order and subset of named tuples

d1bb3b2

oxinabox reviewed Nov 11, 2021

View reviewed changes

mzgubic reviewed Nov 12, 2021

View reviewed changes

src/projection.jl Show resolved Hide resolved

mcabbott added the ProjectTo related to the projection functionality label Nov 24, 2021

oxinabox approved these changes Nov 30, 2021

View reviewed changes

devmotion merged commit addf6d9 into main Dec 5, 2021

devmotion deleted the dw/projectto_namedtuple branch December 5, 2021 23:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `ProjectTo(::NamedTuple)` #515

Add `ProjectTo(::NamedTuple)` #515

devmotion commented Nov 10, 2021

devmotion commented Nov 10, 2021

mcabbott commented Nov 10, 2021

mzgubic commented Nov 11, 2021 •

edited

Loading

devmotion commented Nov 11, 2021

mzgubic commented Nov 11, 2021

devmotion commented Nov 11, 2021

mcabbott commented Nov 11, 2021

devmotion commented Nov 11, 2021

codecov-commenter commented Nov 11, 2021 •

edited

Loading

oxinabox Nov 11, 2021

devmotion Nov 11, 2021

oxinabox Nov 11, 2021

devmotion Nov 11, 2021 •

edited

Loading

mcabbott Nov 11, 2021

oxinabox Nov 11, 2021

mcabbott Nov 11, 2021

mcabbott Dec 3, 2021

oxinabox Dec 3, 2021

devmotion Dec 3, 2021

devmotion Dec 5, 2021

mcabbott Dec 5, 2021

devmotion commented Nov 11, 2021

devmotion commented Nov 27, 2021

oxinabox left a comment

Add ProjectTo(::NamedTuple) #515

Add ProjectTo(::NamedTuple) #515

Conversation

devmotion commented Nov 10, 2021

devmotion commented Nov 10, 2021

mcabbott commented Nov 10, 2021

mzgubic commented Nov 11, 2021 • edited Loading

devmotion commented Nov 11, 2021

mzgubic commented Nov 11, 2021

devmotion commented Nov 11, 2021

mcabbott commented Nov 11, 2021

devmotion commented Nov 11, 2021

codecov-commenter commented Nov 11, 2021 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

devmotion Nov 11, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

devmotion commented Nov 11, 2021

devmotion commented Nov 27, 2021

oxinabox left a comment

Choose a reason for hiding this comment

Add `ProjectTo(::NamedTuple)` #515

Add `ProjectTo(::NamedTuple)` #515

mzgubic commented Nov 11, 2021 •

edited

Loading

codecov-commenter commented Nov 11, 2021 •

edited

Loading

devmotion Nov 11, 2021 •

edited

Loading