Skip to content

RFC: accept Tuple Tangent for arrays? #444

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mcabbott wants to merge 1 commit into from
Closed

RFC: accept Tuple Tangent for arrays? #444

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/projection.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,6 +241,12 @@ function (project::ProjectTo{AbstractArray})(dx::Number) # ... so we restore fro
return fill(project.element(dx))
end

# Accept the Tangent corresponding to a Tuple -- Zygote's splats produce these
function (project::ProjectTo{AbstractArray})(dx::Tangent{<:Any, <:Tuple})

@oxinabox oxinabox Aug 19, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think i prefer this to be written as:

Suggested change
function (project::ProjectTo{AbstractArray})(dx::Tangent{<:Any, <:Tuple})
function (project::ProjectTo{AbstractArray})(dx::Tangent{<:Tuple})

so we don't need to mention how this backed.

@mcabbott mcabbott Aug 19, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think Zygote managed to produce a Tangent{Any, Tuple{...}}, which was why my pirate method specified that, here:

https://github.com/FluxML/Zygote.jl/pull/1044/files#diff-e0bc7da8f1a33a59f5ecfa67257c04038f0b4915b3f74bdf39780818fd0010a2R162

But I was squashing bugs as fast as I could & didn't track it down any further. Can circle back now that things basically pass.

@oxinabox oxinabox Aug 19, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

right, fair enough.
If we have to do it this way that is also fine.
Maybe just a commond is woth adding to clarify that this is the tangent for something with primal type tangent

Zygote does lose primal type information, that is a thing
We might well be able to teach it a bit more about tuples more easily than the general case.

@mcabbott mcabbott Aug 22, 2021

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

BTW, easy examples of this Tangent{Any} problem:

julia> using Zygote

julia> function Zygote.ChainRulesCore.rrule(::typeof(identity), x::AbstractArray)
         x, dx -> (Zygote.NoTangent(), @show dx)
       end

julia> Zygote.gradient(x -> max(identity(x)...), [1,2,3])
dx = Tangent{Any}(0, 0, 1)
((0, 0, 1),)

julia> Zygote.gradient(x -> sum(identity(x).parent), [1,2,3]')
dx = Tangent{Any}(parent = 3-element Fill{Int64}: entries equal to 1,)
((parent = 3-element Fill{Int64}: entries equal to 1,),)

dy = reshape(collect(backing(dx)), project.axes)
return project(dy)
end

# Ref -- works like a zero-array, also allows restoration from a number:
ProjectTo(x::Ref) = ProjectTo{Ref}(; x=ProjectTo(x[]))
(project::ProjectTo{Ref})(dx::Ref) = Ref(project.x(dx[]))
Expand Down