Add rrules for `extrema`, `findmax`, `maximum` by mcabbott · Pull Request #480 · JuliaDiff/ChainRules.jl

mcabbott · 2021-07-27T10:51:11Z

Aims to address FluxML/Zygote.jl#1034 by widening the type of the array of zeros it writes into. And, while there, fixes some related functions:

julia> using Zygote

julia> @btime gradient(maximum, $(rand(1000)));  # old definition
  2.111 μs (18 allocations: 8.36 KiB)

julia> @btime gradient(first∘extrema, $(rand(1000)));  # works, but 3000x slower
  6.402 ms (62198 allocations: 9.75 MiB)

julia> @btime gradient(first∘findmax, $(rand(1000)));
  28.287 ms (204741 allocations: 8.42 MiB)

julia> @btime gradient(x -> sum(maximum(x, dims=1)), $(rand(100, 100)));
  20.250 μs (4 allocations: 80.84 KiB)

julia> @btime gradient(x -> sum(first, extrema(x, dims=1)[1]), $(rand(100, 100)));  # fails
ERROR: Mutating arrays is not supported -- called setindex!(::Matrix{Tuple{Float64, Float64}}, _...)

Still some possible bugs in frules, or their tests?

test/rulesets/Base/array.jl

src/rulesets/Base/array.jl

test/rulesets/Base/array.jl

codecov-commenter · 2021-07-28T04:45:02Z

Codecov Report

Merging #480 (e8f3a67) into master (e6a27db) will decrease coverage by 0.32%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #480      +/-   ##
==========================================
- Coverage   98.38%   98.05%   -0.33%     
==========================================
  Files          21       22       +1     
  Lines        2287     2414     +127     
==========================================
+ Hits         2250     2367     +117     
- Misses         37       47      +10

Impacted Files	Coverage Δ
src/rulesets/Base/nondiff.jl	`100.00% <ø> (+33.33%)`	⬆️
src/rulesets/Base/array.jl	`98.88% <100.00%> (-1.12%)`	⬇️
src/rulesets/LinearAlgebra/structured.jl	`94.00% <0.00%> (-5.30%)`	⬇️
src/rulesets/Base/base.jl	`100.00% <0.00%> (ø)`
src/rulesets/Core/core.jl	`100.00% <0.00%> (ø)`

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 e6a27db...e8f3a67. Read the comment docs.

src/rulesets/Base/array.jl

mcabbott · 2021-07-31T15:31:37Z

test/rulesets/Base/array.jl

+@testset "$findm" for findm in [findmax, findmin]
+    @test_skip test_rrule(findm, rand(10), output_tangent = (rand(), NoTangent()), check_inferred=false)  # error?
+    @test_skip test_rrule(findm, rand(3,4), fkwargs=(dims=1,), output_tangent = (rand(1,4), 999), check_inferred=false) # DimensionMismatch("second dimension of A, 12, does not match length of x, 5"), wtf?
+end
+    @test rrule(findmax, [1,2,33])[1] == (33, 3)
+    @test rrule(findmin, [11,22,33])[1] == (11, 1)
+
+    @test [0,0,1] == @inferred unthunk(rrule(findmax, [1,2,3])[2]((1.0, nothing))[2])
+    @test [1,0,0] == @inferred unthunk(rrule(findmin, [1,2,3])[2]((1.0, nothing))[2])


Not sure why test_rrule fails here, but explicit tests work. The error is:

julia> test_rrule(findm, rand(10), output_tangent = (rand(), NoTangent()), check_inferred=false) test_rrule: findmax on Vector{Float64}: Error During Test at /Users/me/.julia/packages/ChainRulesTestUtils/8380y/src/testers.jl:191 Got exception outside of a @test DimensionMismatch("second dimension of A, 2, does not match length of x, 1") Stacktrace: [1] gemv!(y::Vector{Float64}, tA::Char, A::Matrix{Float64}, x::Vector{Float64}, α::Bool, β::Bool) @ LinearAlgebra ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:477 [2] mul! @ ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:87 [inlined] [3] mul! @ ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:255 [inlined] [4] *(tA::Transpose{Float64, Matrix{Float64}}, x::Vector{Float64}) @ LinearAlgebra ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:80 [5] _j′vp(fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, f::Function, ȳ::Vector{Float64}, x::Vector{Float64}) @ FiniteDifferences ~/.julia/packages/FiniteDifferences/aqPCI/src/grad.jl:80 [6] j′vp(fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, f::Function, ȳ::Tuple{Float64, NoTangent}, x::Vector{Float64}) @ FiniteDifferences ~/.julia/packages/FiniteDifferences/aqPCI/src/grad.jl:73 [7] _make_j′vp_call(fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, f::Function, ȳ::Tuple{Float64, NoTangent}, xs::Tuple{typeof(findmax), Vector{Float64}}, ignores::Tuple{Bool, Bool}) @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/8380y/src/finite_difference_calls.jl:51 [8] macro expansion @ ~/.julia/packages/ChainRulesTestUtils/8380y/src/testers.jl:222 [inlined] [9] macro expansion @ ~/.julia/dev/julia/usr/share/julia/stdlib/v1.8/Test/src/Test.jl:1282 [inlined] [10] test_rrule(config::ChainRulesTestUtils.ADviaRuleConfig, f::typeof(findmax), args::Vector{Float64}; output_tangent::Tuple{Float64, NoTangent}, check_thunked_output_tangent::Bool, fdm::FiniteDifferences.AdaptedFiniteDifferenceMethod{5, 1, FiniteDifferences.UnadaptedFiniteDifferenceMethod{7, 5}}, rrule_f::Function, check_inferred::Bool, fkwargs::NamedTuple{(), Tuple{}}, rtol::Float64, atol::Float64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/8380y/src/testers.jl:194

In revised code, this one id fixed, but errors persist for dims=1 etc. cases

src/rulesets/Base/array.jl

oxinabox · 2021-07-31T19:34:39Z

I am on leave most of next week.
It would be very useful if someone else can pick up the review on this one.

src/rulesets/Base/array.jl

mzgubic · 2021-08-02T13:53:30Z

src/rulesets/Base/array.jl

+    @eval function rrule(::typeof($findm), x::AbstractArray{<:Number}; dims=:)
+        y, ind = $findm(x; dims=dims)
+        project = ProjectTo(x)
+        function $findm_pullback((dy, _))


Should this take a Tangent instead? I think we can still dispatch on dy being an AbstractZero

Maybe? I am a bit confused about Tangent. I was trying things out with Zygote and they appear to work, but perhaps this would still work if the signature was findm_pullback(::Tangent).

Ah, never mind, I thought the destructuring places a constraint. It's fine this way

Added a comment & checked

The method below $findm_pullback(::Tuple{AbstractZero, Any}) will however not accept a Tangent.

Should it be Tangent{<:Any, <: Tuple{AbstractZero, Any}}? Or just dy isa AbstractZero && return (NoTangent(), NoTangent())?

(Now changed to a branch, seems simplest.)

The method below $findm_pullback(::Tuple{AbstractZero, Any}) will however not accept a Tangent.

this is fixed now, right?

Yes, I killed that method, and just check the type after destructuring.

mzgubic · 2021-08-02T14:04:35Z

src/rulesets/Base/array.jl

+        function $findm_pullback((dy, _))
+            x_thunk = @thunk begin
+                dx = fill!(similar(x, eltype(dy)), false)
+                view(dx, ind) .= dy  # possibly 0-dim view, allows dy::Number and dy::Array, and dx::CuArray
+                project(dx)
+            end
+            x_ithunk = InplaceableThunk(x_thunk) do dx
+                view(dx, ind) .= view(dx, ind) .+ dy  # this could be .+=, but not on Julia 1.0
+                dx
+            end
+            return (NoTangent(), x_ithunk)
+        end


could we close over size of x only here?

alternatively, I wonder whether we could reuse the rrule for getindex?
i.e. something like

Suggested change

function $findm_pullback((dy, _))

x_thunk = @thunk begin

dx = fill!(similar(x, eltype(dy)), false)

view(dx, ind) .= dy # possibly 0-dim view, allows dy::Number and dy::Array, and dx::CuArray

project(dx)

end

x_ithunk = InplaceableThunk(x_thunk) do dx

view(dx, ind) .= view(dx, ind) .+ dy # this could be .+=, but not on Julia 1.0

dx

end

return (NoTangent(), x_ithunk)

end

function $findm_pullback((dy, _))

_, getindex_back = rrule(getindex, x, ind)

return getindex_back(dy)[1:2]

end

could we close over size of x only here?

There is similar(typeof(x), size(x)) but no similar(typeof(x), T, size(x)). Of course very often ProjectTo is going to ensure that the eltype should really be the same as x's, but not quite always. It's a bit awkward.

getindex has @thunk(getindex_add!(zero(x))) which seems worse -- it is not always going to be mutable, and it won't handle structure well, e.g. zero(Diagonal([1,2,3]))[1,2] = 9.

Agree that getting it right in one place makes some sense. Zygote now has a special struct for scalar getindex, especially because using that repeatedly in a loop seems common. That does not seem so common for maximum, which could mean the weirdness doesn't pay for itself? Or maybe InplaceableStuff will make that obsolete anyway.

I am not sure I understand the getindex comment. Isn't similar also going to make a Diagonal, just like zero does?

Oh yea, both mess that one up, you need similar(Diagonal(rand(3)), Int, (3,3)) to get something sure to be writeable. It's zero(SA[1,2,3]) which is worse than similar.

Or maybe InplaceableStuff will make that obsolete anyway.

Is that your PR which makes InplaceableThunk take the third argument? I like that idea

Oh maybe that too. But I meant that Zygote.OneElement is part of trying to speed up scalar indexing in a loop. The next step is FluxML/Zygote.jl#981 . But this is a bit of a Zygote-style hack, and the eventual version involve ChainRules's in-place stuff. Then a tight loop might generate a million thunks, instead of a million OneElement "arrays".

src/rulesets/Base/array.jl

mzgubic

generally LGTM, should we just move the view and sum to ChainRulesCore first?

JuliaDiff/ChainRulesCore.jl#425

mcabbott · 2021-08-03T20:30:00Z

Thanks, will update to use new CRCore.

Check that second derivatives aren't completely wrong (with Zygote's rule disabled):

julia> Zygote.gradient(x -> maximum(Zygote.gradient(y -> sum(abs2, maximum(y, dims=1)), x)[1]), [1 2; 3 4])
([0.0 0.0; 0.0 2.0],)
# was ERROR: Mutating arrays is not supported on Zygote v0.6.17

julia> ForwardDiff.gradient(x -> maximum(ForwardDiff.gradient(y -> sum(abs2, maximum(y, dims=1)), x)), [1 2; 3 4])
2×2 Matrix{Int64}:
 0  0
 0  2

mcabbott · 2021-08-03T20:45:54Z

test/rulesets/Base/array.jl

+    @test_skip test_frule(findmin, rand(3,4)) # StackOverflowError, CartesianIndex{2}(index::Tuple{Float64, Float64}) (repeats 79984 times) & TypeError: in new, expected Tuple{Int64, Int64}, got a value of type Tuple{Float64, Float64}
+    @test_skip test_frule(findmin, rand(3,4), output_tangent = (rand(), NoTangent()))
+    @test_skip test_frule(findmin, rand(3,4), fkwargs=(dims=1,))
+
+    # Reverse
+    test_rrule(findmin, rand(10), output_tangent = (rand(), false))
+    test_rrule(findmax, rand(10), output_tangent = (rand(), false))
+    @test [0 0; 0 5] == @inferred unthunk(rrule(findmax, [1 2; 3 4])[2]((5.0, nothing))[2])
+    @test_skip test_rrule(findmin, rand(10), output_tangent = (rand(), NoTangent()), check_inferred=false)  # DimensionMismatch from FiniteDifferences
+    @test_skip test_rrule(findmax, rand(5,3), output_tangent = (rand(), false), check_inferred=false)  # DimensionMismatch from FiniteDifferences


Was still a bit unhappy about these tests.

I think this would solve it: JuliaDiff/FiniteDifferences.jl#188

Any thoughts on this? Maybe should merge without these tests.

I would suggest keeping as it is and adding a comment that a lot of dimension mismatches would be solved by fixing the JuliaDiff/FiniteDifferences.jl#188

FWIW the current error here is:

julia> test_frule(findmin, rand(3,4)) test_frule: findmin on Matrix{Float64}: Error During Test at /Users/me/.julia/packages/ChainRulesTestUtils/73Y9Q/src/testers.jl:118 Got exception outside of a @test iteration is deliberately unsupported for CartesianIndex. Use `I` rather than `I...`, or use `Tuple(I)...` Stacktrace: [1] error(s::String) @ Base ./error.jl:33 [2] iterate(#unused#::CartesianIndex{2}) @ Base.IteratorsMD ./multidimensional.jl:167 [3] copyto!(dest::Vector{Int64}, src::CartesianIndex{2}) @ Base ./abstractarray.jl:901 [4] _collect(cont::UnitRange{Int64}, itr::CartesianIndex{2}, #unused#::Base.HasEltype, isz::Base.HasLength) @ Base ./array.jl:715 [5] collect(itr::CartesianIndex{2}) @ Base ./array.jl:709 [6] test_approx(actual::CartesianIndex{2}, expected::CartesianIndex{2}, msg::Any; kwargs::Base.Pairs{Symbol, V, Tuple{Vararg{Symbol, N}}, NamedTuple{names, T}} where {V, N, names, T<:Tuple{Vararg{Any, N}}}) @ ChainRulesTestUtils ~/.julia/packages/ChainRulesTestUtils/73Y9Q/src/check_result.jl:141

mzgubic · 2021-11-24T10:25:19Z

test/rulesets/Base/array.jl

+    # Reverse with dims:
+    @test [0 0; 5 6] == @inferred unthunk(rrule(findmax, [1 2; 3 4], dims=1)[2](([5 6], nothing))[2])
+    @test [5 0; 6 0] == @inferred unthunk(rrule(findmin, [1 2; 3 4], dims=2)[2]((hcat([5,6]), nothing))[2])
+    @test_skip test_rrule(findmin, rand(3,4), fkwargs=(dims=1,), output_tangent = (rand(1,4), NoTangent()), check_inferred=false) # DimensionMismatch("second dimension of A, 12, does not match length of x, 5"), wtf?


These could be solved by FiniteDifferences knowing that CartesianIndex is not perturbable: JuliaDiff/FiniteDifferences.jl#196

Here's a PR to fix it: JuliaDiff/FiniteDifferences.jl#197

that's been merged, it needs FiniteDifferences 0.12.20

That now passes, thanks!

mcabbott · 2021-11-24T15:48:11Z

Testsets pass locally, will merge when (if!) CI agrees.

Besides tests, upgraded today to allow arrays of arrays.

There are a few tests skipped, I think due to weird FiniteDifferences errors. But I think the rules work, and e.g. the findmax rules are in fact tested via the maximum rules.

mcabbott commented Jul 27, 2021

View reviewed changes

test/rulesets/Base/array.jl Outdated Show resolved Hide resolved

mcabbott commented Jul 27, 2021

View reviewed changes

src/rulesets/Base/array.jl Outdated Show resolved Hide resolved

mcabbott commented Jul 27, 2021

View reviewed changes

test/rulesets/Base/array.jl Outdated Show resolved Hide resolved

mcabbott mentioned this pull request Jul 27, 2021

Document subgradient convention JuliaDiff/ChainRulesCore.jl#404

Closed

mcabbott commented Jul 28, 2021

View reviewed changes

src/rulesets/Base/array.jl Outdated Show resolved Hide resolved

mcabbott mentioned this pull request Jul 30, 2021

Algorithms which branch to simpler implementations EnzymeAD/Enzyme.jl#114

Open

mcabbott commented Jul 31, 2021

View reviewed changes

mcabbott marked this pull request as ready for review July 31, 2021 15:50

mcabbott commented Jul 31, 2021

View reviewed changes

src/rulesets/Base/array.jl Outdated Show resolved Hide resolved

oxinabox reviewed Jul 31, 2021

View reviewed changes

src/rulesets/Base/array.jl Outdated Show resolved Hide resolved

mzgubic reviewed Aug 2, 2021

View reviewed changes

mzgubic approved these changes Aug 3, 2021

View reviewed changes

mcabbott force-pushed the extrema branch from 87db7e9 to 46571cf Compare August 3, 2021 20:35

mcabbott commented Aug 3, 2021

View reviewed changes

mcabbott mentioned this pull request Aug 5, 2021

Configured rule for maximum(f, xs) #490

Closed

mzgubic reviewed Nov 24, 2021

View reviewed changes

mcabbott added 11 commits November 24, 2021 09:03

rules for extrema, findmax, maximum

fc6fad9

fixup extrema

0e01222

symmetric maximum rule

4962e4a

promote types by hand

a60d244

argmax?

457a3f2

allow more zeros

98df2f1

upgrade tests

2e63a86

don't do symmetric convention

47d9822

tests

8eb4b1c

fix 1.0

2df9681

rm symmetric versions

9fdfec3

mcabbott added 18 commits November 24, 2021 09:03

move extrema to last

bf0bf2e

tidy

5c57f44

fixup extrema

1ba053a

tests

15ccc56

tests

7ef9a32

use eval loop, tidy, tests

3af83c0

forward rules for maximum

8fca98a

frules for findmax

ee7d97f

tidy

67ea14d

widen similar to ensure writeability

b17be68

comments

8d96687

dispatch -> branch

5bd06ad

allow for second derivatives

9813783

frule?

4039c93

update to use CRC 1.3

8e07ada

better writezero?

1722ca5

fix tests

656602a

allow arrays of arrays

77d7526

mcabbott force-pushed the extrema branch from 46571cf to 77d7526 Compare November 24, 2021 15:31

version

b17d91e

mcabbott merged commit 605354c into JuliaDiff:main Nov 24, 2021

mcabbott deleted the extrema branch November 24, 2021 16:35

Conversation

mcabbott commented Jul 27, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov-commenter commented Jul 28, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Uh oh!

mcabbott Jul 31, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

oxinabox commented Jul 31, 2021

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mcabbott Aug 2, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mcabbott Aug 2, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

mzgubic left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mcabbott commented Aug 3, 2021

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mzgubic Nov 24, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mcabbott commented Jul 27, 2021 •

edited

Loading

codecov-commenter commented Jul 28, 2021 •

edited

Loading

mcabbott Jul 31, 2021 •

edited

Loading

mcabbott Aug 2, 2021 •

edited

Loading

mcabbott Aug 2, 2021 •

edited

Loading

mzgubic left a comment •

edited

Loading

mzgubic Nov 24, 2021 •

edited

Loading