`ProjectTo{<:Tangent}` for tuples & Ref #488

mcabbott · 2021-10-11T01:33:10Z

Replaces #457

codecov-commenter · 2021-10-11T01:35:28Z

Codecov Report

Merging #488 (3e5cda1) into main (d0c3599) will decrease coverage by 0.06%.
The diff coverage is 92.30%.

@@            Coverage Diff             @@
##             main     #488      +/-   ##
==========================================
- Coverage   93.00%   92.94%   -0.07%     
==========================================
  Files          15       15              
  Lines         801      822      +21     
==========================================
+ Hits          745      764      +19     
- Misses         56       58       +2

Impacted Files	Coverage Δ
src/projection.jl	`97.47% <92.30%> (-0.68%)`	⬇️

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 d0c3599...3e5cda1. Read the comment docs.

mcabbott · 2021-10-11T01:35:40Z

src/projection.jl

-# Since this works like a zero-array in broadcasting, it should also accept a number:
-(project::ProjectTo{Ref})(dx::Number) = Tangent{project.type}(; x=project.x(dx))
+
+(project::ProjectTo{<:Tangent{<:Tuple}})(dx::Tangent) = project(backing(dx))


This means that projection on the output of this projector will disassemble the Tangent and re-process the Tuple inside. I'm not sure that's ideal. Maybe it's safe to pass on all Tangents without further investigation?

I think this should constant-fold out for most cases we care about.
Because it will see the ProjectTo for the backing elements as already being the right type, which is known from the type in the project.elements

Do you want to check some with @code_typed / Cthulu ?

It is not safe to pass on all Tangents, because the tangent could be wrapping Complex Number/ Dense array that we need to fix.

What I wonder is whether I can think of the "mathematical" steps involving arrays etc. as being separate from the "structural" steps involving Tangents. If the first project, and then the backward flow assembles and de-assembles a Tangent, can this Tangent have "crossed a boundary" such that it belongs to a different argument type and hence may need further projection? I mostly think it can't; it would have to get un-packaed and those pieces operated on. But I'm not very sure.

I see.
Yeah, isn't this kind of true for many many operations that project?
They wouldn't need to project if were sure that the "mathematical" step before had projected and so only given them something good?
(also applies for if the mathematical step was a human constructing the tangents)

But it gets fuzzy around the edges?
How would sum(sum, ((Diagonal[1f0,2f0]), (Diagonal[1f0,2f0]))) go down?

oxinabox · 2021-10-11T14:02:09Z

The formatter is unhappy see: https://github.com/JuliaDiff/ChainRulesCore.jl/pull/488/checks?check_run_id=3854350212#step:5:187

I think it wasn't allowed to post suggestions as it was made from a fork.
(See #489)

That should be fixed, or the formatter will complain in other PRs.
(and also it is i guess right)

src/projection.jl

oxinabox · 2021-10-11T14:28:43Z

src/projection.jl

+    dy = reshape(dx, axes(project.elements))  # allows for dx::OffsetArray
+    dz = ntuple(i -> project.elements[i](dy[i]), length(project.elements))
+    return project_type(project)(dz...)
+end


The above Tuple and AbstractArray cases are just optimizations of a general iterator one:

function (project::ProjectTo{<:Tangent{<:Tuple}})(dxs) # iterator fallback dzs = (f(dx) for (f, dx) in zip(project.elements, dxs)) return project_type(project)(dzs...) end

Should we have that as well?
And then we can note the others as just being optimizations?

We could do. Do you have something in mind which might produce some weird type?

If some NamedTuple leaks from Zygote, I think this will produce stranger error messages, since it may make a Tangent of the wrong length?

I don't really have it in mind, more is that that is the general case we are handling.
It is weird that we only actually handle the two optimizable versions?

I guess I'm not entirely sure we need the array version at all. I was thinking about things like broadcasting, although that handles it explicitly... but map doesn't:

julia> Zygote.pullback(x -> sum(map(+, x, [1,2])), (1,2))[2](1) ([1, 1],) julia> gradient(x -> sum(map(+, x, [1,2])), (1,2)) # uses projection ((1.0, 1.0),)

I would be down with seeing it removed til we know we need it.

My vote is to keep this, although I can't think of another example besides map right now which uses this. It looks like gradient(x -> sum([i^2 for i in x]), (1,2)) does not.

src/projection.jl

oxinabox · 2021-10-11T14:33:22Z

src/projection.jl

+    if length(dx) != len
+        str = "tuple with length(x) == $len cannot have a gradient with length(dx) == $(length(dx))"
+        throw(DimensionMismatch(str))
+    end


will this not be caught by map?

If we removed this check then this would basically be the general iterator fallback case.
https://github.com/JuliaDiff/ChainRulesCore.jl/pull/488/files#r726169970

It will, but the error is much less friendly... and might be a bug, JuliaLang/julia#42216

src/projection.jl

oxinabox

Basically LGTM.
A few things to address.
Once sorted merge when happy

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

mcabbott · 2021-10-15T03:55:15Z

Ref is I think the only mutable type for which we have defined projection rules. I think the lesson of JuliaDiff/ChainRules.jl#539 (comment) is that Zygote's structural tangents are weird for such things.

Does Tangent make any such distinction, which needs handling here?

If not, then it is I think entirely a matter of fixing the Zygote/ChainRules interface. Maybe wrap_chainrules_input should handle this; it had a method for NamedTuples but not for Ref. There appear to be zero tests. In fact I'm pretty surprised I didn't stumble across this while fiddling with projection rules. Makes me wonder a bit if we ought to have more explicit downstream testing, i.e. don't just trust the package's tests, but add tests here (or in ChainRules) using packages explicitly, while thinking about any given feature.

oxinabox · 2021-10-15T10:31:31Z

Does Tangent make any such distinction, which needs handling here?

It does not.
There is an open issue for it, as it can come up in forward-mode.
#105

If not, then it is I think entirely a matter of fixing the Zygote/ChainRules interface. Maybe wrap_chainrules_input should handle this; it had a method for NamedTuples but not for Ref.

Correct.
When I wrote that I did not know Zygote did this.

There appear to be zero tests. In fact I'm pretty surprised I didn't stumble across this while fiddling with projection rules.

You might be right, there are not tests of it directly.
I think in most cases though there are tests that hit it implictly.
But not for the mutable struct case I guess.
(The direct tests are mostly focused around making sure that the rules are picked up, and around things that were wrong before like functions that return tuples and functions that take keyword arguments)

Makes me wonder a bit if we ought to have more explicit downstream testing, i.e. don't just trust the package's tests, but add tests here (or in ChainRules) using packages explicitly, while thinking about any given feature.

I have worked on packages that have direct downstream tests.
It is basically hell, due to the near circular dependency.
It causes all kinds of problems with version bounds and things get fragile and start to display "intimate knowledge" of each other.
Much better is the way we are doing it here, which includes the downstream tests automatically passing if compat is not permitted.
We can just fix the tests in Zygote etc.

We should test features in ways that are like how they are intended to be used by downstream packages.
But we should not have tests that can't be fixed by changing this package.

mcabbott · 2021-10-15T11:51:38Z

Ok, thanks!

It looks like we need this PR first. Then the fix is one line:

julia> gradient(x -> (x.x)^2, Ref(3))
ERROR: MethodError: no method matching (::ProjectTo{Float64, NamedTuple{(), Tuple{}}})(::NamedTuple{(:x,), Tuple{Int64}})
  ...
Stacktrace:
 [1] (::ProjectTo{Tangent{Base.RefValue{Int64}}, NamedTuple{(:x,), Tuple{ProjectTo{Float64, NamedTuple{(), Tuple{}}}}}})(dx::Base.RefValue{Any})
   @ ChainRulesCore ~/.julia/dev/ChainRulesCore/src/projection.jl:282
 [2] _project
...

julia> @eval Zygote wrap_chainrules_input(x::Ref) = wrap_chainrules_input(x[]);

julia> gradient(x -> (x.x)^2, Ref(3))
((x = 6.0,),)

What that one line won't do is allow wrap_chainrules_output to restore this Ref(NamedTuple) thing. I don't see a way to do that without #105. Maybe I need to think up a way to test if that matters to Zygote.

mcabbott · 2021-10-15T11:55:07Z

hell, due to the near circular dependency.

Point taken. It might be worthwhile developing a habit of writing a few Zygote tests while working on something here, to contribute there. Perhaps all gathered in one file, ideally to be coped to Diffractor later?

devmotion · 2021-11-09T22:08:49Z

@mcabbott Did you on purpose not add any definitions for ProjectTo(x::NamedTuple)? They would be needed in JuliaDiff/ChainRulesTestUtils.jl#224, I tried to explain it a bit in JuliaDiff/ChainRulesTestUtils.jl#224 (comment).

mcabbott · 2021-11-09T22:24:50Z

No, I planned to but cut it out of this PR in the end. It ought to exist though, someone just has to write it.

mcabbott · 2021-11-09T22:48:49Z

The simplest version is something like:

function ProjectTo(x::NamedTuple)
    elements = map(ProjectTo, x)
    if Tuple(elements) isa NTuple{<:Any,ProjectTo{<:AbstractZero}}
        return ProjectTo{NoTangent}()
    else
        return ProjectTo{Tangent{typeof(x)}}(; elements...)
    end
end
(project::ProjectTo{<:Tangent{<:NamedTuple}})(dx::Tangent) = project(backing(dx))
function (project::ProjectTo{<:Tangent{<:NamedTuple}})(dx::NamedTuple)
    dy = map((f, x) -> f(x), backing(project), dx)
    return project_type(project)(; dy...)
end

That demands exact equality of the names, which I think is what you want.

tuple -> tangent, take 1

17a941e

mcabbott commented Oct 11, 2021

View reviewed changes

oxinabox reviewed Oct 11, 2021

View reviewed changes

src/projection.jl Outdated Show resolved Hide resolved

mcabbott mentioned this pull request Oct 11, 2021

No method matching for mapreduce? JuliaDiff/ChainRules.jl#539

Closed

oxinabox reviewed Oct 11, 2021

View reviewed changes

src/projection.jl Outdated Show resolved Hide resolved

oxinabox reviewed Oct 11, 2021

View reviewed changes

src/projection.jl Outdated Show resolved Hide resolved

oxinabox reviewed Oct 11, 2021

View reviewed changes

src/projection.jl Outdated Show resolved Hide resolved

oxinabox approved these changes Oct 11, 2021

View reviewed changes

mcabbott and others added 3 commits October 11, 2021 10:38

Apply 4 suggestions

8b8314e

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

test + version

1d75133

rm "only" for 1.0

ecd318f

mcabbott added 2 commits October 15, 2021 08:16

version + comments

15756a4

Merge branch 'main' into tangent2

3e5cda1

mcabbott merged commit 38c4089 into JuliaDiff:main Oct 15, 2021

mcabbott deleted the tangent2 branch October 15, 2021 13:36

oxinabox mentioned this pull request Oct 15, 2021

Make wrap_chainrules_input understand mutable structs FluxML/Zygote.jl#1102

Merged

mcabbott mentioned this pull request Oct 16, 2021

ProjectTo called on tuples #440

Closed

mzgubic mentioned this pull request Nov 10, 2021

implement ProjectTo(::NamedTuple) #511

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

`ProjectTo{<:Tangent}` for tuples & Ref #488

`ProjectTo{<:Tangent}` for tuples & Ref #488

mcabbott commented Oct 11, 2021

codecov-commenter commented Oct 11, 2021 •

edited

Loading

mcabbott Oct 11, 2021

oxinabox Oct 11, 2021

mcabbott Oct 11, 2021

oxinabox Oct 11, 2021

oxinabox commented Oct 11, 2021

oxinabox Oct 11, 2021 •

edited

Loading

mcabbott Oct 11, 2021

oxinabox Oct 11, 2021

mcabbott Oct 11, 2021

oxinabox Oct 11, 2021

mcabbott Oct 15, 2021

oxinabox Oct 11, 2021 •

edited

Loading

mcabbott Oct 11, 2021

oxinabox left a comment

mcabbott commented Oct 15, 2021

oxinabox commented Oct 15, 2021

mcabbott commented Oct 15, 2021

mcabbott commented Oct 15, 2021

devmotion commented Nov 9, 2021

mcabbott commented Nov 9, 2021 •

edited

Loading

mcabbott commented Nov 9, 2021

ProjectTo{<:Tangent} for tuples & Ref #488

ProjectTo{<:Tangent} for tuples & Ref #488

Conversation

mcabbott commented Oct 11, 2021

codecov-commenter commented Oct 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

Choose a reason for hiding this comment

oxinabox commented Oct 11, 2021

oxinabox Oct 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

oxinabox Oct 11, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

oxinabox left a comment

Choose a reason for hiding this comment

mcabbott commented Oct 15, 2021

oxinabox commented Oct 15, 2021

mcabbott commented Oct 15, 2021

mcabbott commented Oct 15, 2021

devmotion commented Nov 9, 2021

mcabbott commented Nov 9, 2021 • edited Loading

mcabbott commented Nov 9, 2021

`ProjectTo{<:Tangent}` for tuples & Ref #488

`ProjectTo{<:Tangent}` for tuples & Ref #488

codecov-commenter commented Oct 11, 2021 •

edited

Loading

oxinabox Oct 11, 2021 •

edited

Loading

oxinabox Oct 11, 2021 •

edited

Loading

mcabbott commented Nov 9, 2021 •

edited

Loading