Simplest prod(x; dims) gradient

[mcabbott]

The current gradient for prod(x; dims) gives incorrect results, this PR fixes it (parallel to FluxML/Tracker.jl#1 ):

julia> using Zygote, ForwardDiff

julia> r = rand(2,3,2);

julia> ForwardDiff.gradient(w->sum(prod(w, dims=(2,3))), r)
2×3×2 Array{Float64,3}:
[:, :, 1] =
 0.00131643  0.000954347  0.0051387 
 0.0177437   0.0354628    0.00934587

[:, :, 2] =
 0.00434307  0.0140455   0.00152818
 0.0151417   0.00464615  0.00451601

julia> Zygote.gradient(w->sum(prod(w, dims=(2,3))), r)[1] # wrong answer!
2×3×2 Array{Float64,3}:
[:, :, 1] =
 5.93867e-6  4.30525e-6  2.31817e-5
 1.60301e-5  3.2038e-5   8.44331e-6

[:, :, 2] =
 1.95925e-5  6.33622e-5  6.89391e-6
 1.36795e-5  4.19746e-6  4.07989e-6

julia> Zygote.@adjoint function prod(xs; dims = :) # as in this PR
         p = prod(xs; dims = dims)
         p, Δ -> (p ./ xs .* Δ,)
       end

julia> Zygote.refresh()

julia> Zygote.gradient(w->sum(prod(w, dims=(2,3))), r)[1] # now matches ForwardDiff
2×3×2 Array{Float64,3}:
[:, :, 1] =
 0.00131643  0.000954347  0.0051387 
 0.0177437   0.0354628    0.00934587

[:, :, 2] =
 0.00434307  0.0140455   0.00152818
 0.0151417   0.00464615  0.00451601

This does not handle zeros in the array correctly -- see FluxML/Flux.jl#524 for attempts to do that. The circshift(... operation deleted here was a correct (but slow) gradient for prod(x), but is clearly independent of dims.

The example above is almost the same as the one in the tests, which strangely passes, without this PR. Perhaps something is wrong with gradtest?

julia> @test gradtest(x -> prod(x, dims = (2, 3)), (3,4,5))
Test Passed

julia> @test gradtest(x -> prod(x), (3,4,5))
Test Passed

[mcabbott]

Bump?

Just cleaned up the mess that github's merge tool sometimes creates with unicode y, ȳ -> etc. Locally I had test failures of gradtest((x, w) -> conv(x, w, cdims), x, w) etc, which I think are unrelated.

[DhairyaLGandhi]

With a precursory look, it looks about right to me. Hopping on a flight now, but in the meantime

bors try

[bors]

try

Build succeeded

• ci/gitlab/gitlab.com

[CarloLucibello]

bors r+

[CarloLucibello]

@maleadt is bors down?

[DhairyaLGandhi]

bors r+

[maleadt]

https://app.bors.tech/
Looks down.

[CarloLucibello]

bors r+

[bors]

🔒 Permission denied

Existing reviewers: click here to make CarloLucibello a reviewer

[CarloLucibello]

@MikeInnes @dhairyagandhi96 @maleadt could someone add me to bors' reviewers?

[DhairyaLGandhi]

Could you try now?

[CarloLucibello]

bors r+

[bors]

🔒 Permission denied

Existing reviewers: click here to make CarloLucibello a reviewer

[CarloLucibello]

Could you try now?

no luck

[DhairyaLGandhi]

bors delegate=CarloLucibello

[bors]

✌️ CarloLucibello can now approve this pull request. To approve and merge a pull request, simply reply with bors r+. More detailed instructions are available here.

[DhairyaLGandhi]

Should do for now, will look at the dashboard in my morn

[DhairyaLGandhi]

bors r+

[CarloLucibello]

Should do for now, will look at the dashboard in my morn

I think I am in Flux's list now but not in Zygote's

[DhairyaLGandhi]

That's what caught me off guard, bors-ng/bors-ng#517 suggests the fix too

[bors]

Build succeeded

• ci/gitlab/gitlab.com