update, tidy

Author: mcabbott

src/rulesets/Base/mapreduce.jl

@@ -188,48 +188,49 @@ end
 rrule(::typeof(cumsum), x::AbstractVector) = rrule(cumsum, x; dims=1)
 
 #####
-##### `maximum`, `minimum`
+##### `maximum(f, xs)`, `minimum(f, xs)`
 #####
 
+# Rules for `maximum(x)` live with `findmax(x)` in array.jl
+
 for mimum in (:minimum, :maximum)
-    pullback1 = Symbol(mimum, :_pullback_f)
-    pullback2 = Symbol(mimum, :_pullback_composed)
     findm = Symbol(:find, string(mimum)[1:3])
 
     @eval function rrule(
-        config::RuleConfig{>:HasReverseMode}, ::typeof($mimum), f::F, xs::AbstractArray{<:Number}; dims=:
+        config::RuleConfig{>:HasReverseMode},
+        ::typeof($mimum),
+        f::F,
+        xs::AbstractArray{<:Number};
+        dims = :,
     ) where {F}
         project = ProjectTo(xs)
-
-        # The easy case is when we can use `findmax` to get index, and write into it:
-        if dims isa Colon && VERSION >= v"1.7-"
+        if dims isa Colon && VERSION >= v"1.7"
+            # The easy case is when we can use `findmax` to get index, and write into it:
             y, ind = $findm(f, xs)
-            function $pullback1(dy)
-                # Notice this evaluates `f` one more time, but this shouldn't matter
-                # unless `f` is sateful, in which case both this and `maximum(f.(xs))` 
-                # give undefined results.
+            function minormax_f_back1(dy)
+                # Notice this evaluates `f` one more time, but this shouldn't matter unless `f` is
+                # sateful, in which case both this and `maximum(f.(xs))` give uncertain results.
                 _, one_back = rrule_via_ad(config, f, xs[ind])
-                df, one_dx_raw = one_back(unthunk(dy))
-                one_dx = unthunk(one_dx_raw)
-                x_thunk = @thunk project(_zerolike_writeat(xs, one_dx, dims, ind))
+                df, one_dx = one_back(unthunk(dy))
+                x_thunk = @thunk project(_zerolike_writeat(xs, unthunk(one_dx), dims, ind))
                 x_ithunk = InplaceableThunk(x_thunk) do dxs
-                    view(dxs, ind) .+= one_dx
+                    view(dxs, ind) .+= unthunk(one_dx)  # TODO make _zerolike_writeat handle thunks
                     dxs
                 end
                 return (NoTangent(), df, x_ithunk)
             end
-            return y, $pullback1
+            return y, minormax_f_back1
 
-        # Otherwise, the best path is to broadcast, `maximum(f.(xs); dims)`:
         else
-            mid, cast_back = rrule_via_ad(config, broadcast, f, xs; dims=dims)
-            y, max_back = rrule($mimum, fxs; dims=dims)
-            function $pullback2(dys)
-                _, dmid = max_back(dys)
-                _, df, dxs = cast_back(dmid)  # if cast_back from rrule_via_ad makes an InplaceableThunk,
-                return (NoTangent(), df, project(dxs))  # then this project() will give an error.
+            # Otherwise, the best path is to broadcast, `maximum(f.(xs); dims)`:
+            fxs, cast_back = rrule_via_ad(config, broadcast, f, xs)
+            y, mm_back = rrule($mimum, fxs; dims)
+            function minormax_f_back2(dy)
+                _, dmid = mm_back(dy)
+                _, df, dxs = cast_back(dmid)
+                return (NoTangent(), df, project(dxs))
             end
-            return y, $pullback2
+            return y, minormax_f_back2
         end
 
     end # @eval function rrule(...)

test/rulesets/Base/mapreduce.jl

@@ -120,22 +120,21 @@ const CFG = ChainRulesTestUtils.ADviaRuleConfig()
         @test rrule(SumRuleConfig(), Base.sum, xs, weights) isa Nothing
     end
 
-    @testset "maximum(f, xs)" begin
-        # This calls back into AD
+    VERSION >= v"1.7" && @testset "maximum(f, xs)" begin
         test_rrule(maximum, abs, [-4.0, 2.0, 2.0])
         test_rrule(minimum, sqrt, Float64[1 2; 3 4])
         test_rrule(maximum, Multiplier(2.0), [2.0, 4.0, 8.0])  # Multiplier defined in test_helpers.jl
 
-        # dims keyword -- these need to call `rrule_via_ad(broadcast, ...`
-        @test_skip test_rrule(maximum, sqrt, Float64[1 2; 3 4], fkwargs=(;dims=1), check_inferred=false)
-        @test_skip test_rrule(minimum, abs, randn(3,3), fkwargs=(;dims=2), check_inferred=false)
-
         # repeated -- can't use FiniteDifferences
-        y1, bk1 = rrule(TestConfigReverse(), maximum, abs, [-4.0, 2.0, 4.0, 2.0])  # TestConfigReverse defined in test_helpers.jl
+        y1, bk1 = rrule(CFG, maximum, abs, [-4.0, 2.0, 4.0, 2.0])
         @test y1 === 4.0
         @test unthunk(bk1(10.0)[3]) == [-10, 0, 0, 0]
 
-        # y2, bk2 = rrule(TestConfigReverse(), minimum, abs, [1 2 3; -5 -4 -4], dims=2)
+        # dims keyword -- these need to call `rrule_via_ad(broadcast, ...`, which needs AD
+        @test_skip test_rrule(maximum, sqrt, Float64[1 2; 3 4], fkwargs=(; dims = 1), check_inferred=false)
+        @test_skip test_rrule(minimum, abs, randn(3,3), fkwargs=(; dims = 2), check_inferred=false)
+
+        @test_skip y2, bk2 = rrule(CFG, minimum, abs, [1 2 3; -5 -4 -4], dims = 2)
         # @test y2 == hcat([1, 4])
         # @test unthunk(bk2(hcat([10, 20]))[3]) == [10 0 0; 0 -20 0]
     end

test/test_helpers.jl

@@ -75,25 +75,6 @@ function ChainRulesCore.rrule(::typeof(make_two_vec), x)
     return make_two_vec(x), make_two_vec_pullback
 end
 
-# Trivial rule configurations, allowing `rrule_via_ad` with simple functions:
-struct TestConfigReverse <: RuleConfig{HasReverseMode} end
-function ChainRulesCore.rrule_via_ad(::TestConfigReverse, args...; kw...)
-    if hasmethod(rrule, typeof(args), keys(kw))
-        rrule(args...; kw...)
-    else
-        error("TestConfigReverse can only handle `rrule_via_ad(f, args...)` when there is an rrule method")
-    end
-end
-
-struct TestConfigForwards <: RuleConfig{HasForwardsMode} end
-function ChainRulesCore.frule_via_ad(::TestConfigReverse, args...; kw...)
-    if hasmethod(frule, typeof(args), keys(kw))
-        frule(args...; kw...)
-    else
-        error("TestConfigForwards can only handle `frule_via_ad(f, args...)` when there is an frule method")
-    end
-end
-
 @testset "test_helpers.jl" begin
 
     @testset "Multiplier" begin