Merge pull request #619 from JuliaDiff/revert-615-mean_f_x

Revert "Rule for `mean(f,x)`"

Author: rofinn

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.32.0"
+version = "1.32.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Statistics/statistics.jl

@@ -6,31 +6,19 @@ _denom(x, dims::Colon) = length(x)
 _denom(x, dims::Integer) = size(x, dims)
 _denom(x, dims) = mapreduce(i->size(x, i), Base.mul_prod, unique(dims), init=1)
 
-function rrule(::typeof(mean), x::AbstractArray{<:Union{Real,Complex,AbstractArray}}; dims=:)
-    y_sum, sum_pullback = rrule(sum, x; dims)
+# TODO: We have `mean(f, x; dims)` as of 1.3.0-DEV.36
+# https://github.com/JuliaDiff/ChainRules.jl/issues/85
+function rrule(::typeof(mean), x::AbstractArray{<:Real}; dims=:)
+    y_sum, sum_pullback = rrule(sum, x; dims=dims)
     n = _denom(x, dims)
     function mean_pullback(ȳ)
-        _, ∂x = sum_pullback(unthunk(ȳ) / n)
+        _, ∂sum_x = sum_pullback(ȳ)
+        ∂x = unthunk(∂sum_x) / n
         return (NoTangent(), ∂x)
     end
     return y_sum / n, mean_pullback
 end
 
-function rrule(
-    config::RuleConfig{>:HasReverseMode},
-    ::typeof(mean),
-    f::F,
-    x::AbstractArray{T};
-    dims=:,
-) where {F, T<:Union{Real,Complex,AbstractArray}}
-    y_sum, sum_pullback = rrule(config, sum, f, x; dims)
-    n = _denom(x, dims)
-    function mean_pullback_f(ȳ)
-        return sum_pullback(unthunk(ȳ) / n)
-    end
-    return y_sum / n, mean_pullback_f
-end
-
 #####
 ##### variance
 #####

test/rulesets/Statistics/statistics.jl

@@ -1,32 +1,11 @@
 @testset "mean" begin
-    @testset "mean(x)" begin
-        test_rrule(mean, randn(9))
-        test_rrule(mean, randn(ComplexF64,2,4))
-        test_rrule(mean, transpose(rand(3)))
-        test_rrule(mean, [rand(3) for _ in 1:4]; check_inferred=false)
+    n = 9
+    @testset "Basic" begin
+        test_rrule(mean, randn(n))
     end
     @testset "with dims kwargs" begin
-        test_rrule(mean, randn(9); fkwargs=(;dims=1))
-        test_rrule(mean, randn(9,4); fkwargs=(;dims=2))
-        test_rrule(mean, [rand(2) for _ in 1:3, _ in 1:4]; fkwargs=(;dims=2), check_inferred=false)
-    end
-    @testset "mean(f, x)" begin
-        # This shares its implementation with sum(f, x). Similar tests should cover all cases:
-        test_rrule(mean, abs, [-4.0, 2.0, 2.0])
-        test_rrule(mean, log, rand(3, 4) .+ 1)
-        test_rrule(mean, cbrt, randn(5))
-        test_rrule(mean, Multiplier(2.0), [2.0, 4.0, 8.0])  # defined in test_helpers.jl
-        test_rrule(mean, Divider(1 + rand()), randn(5)) 
-
-        test_rrule(mean, sum, [[2.0, 4.0], [4.0,1.9]]; check_inferred=false)
-
-        test_rrule(mean, log, rand(ComplexF64, 5))
-        test_rrule(mean, sqrt, rand(ComplexF64, 5))
-        test_rrule(mean, abs, rand(ComplexF64, 3, 4))
-        
-        test_rrule(mean, abs, [-2.0 4.0; 5.0 1.9]; fkwargs=(;dims=1))
-        test_rrule(mean, abs, [-2.0 4.0; 5.0 1.9]; fkwargs=(;dims=2))
-        test_rrule(mean, sqrt, rand(ComplexF64, 3, 4); fkwargs=(;dims=(1,)))
+        test_rrule(mean, randn(n); fkwargs=(;dims=1))
+        test_rrule(mean, randn(n,4); fkwargs=(;dims=2))
     end
 end
 

test/test_helpers.jl

@@ -35,28 +35,6 @@ function ChainRulesCore.rrule(m::Multiplier, y, z)
     return m(y, z), Multiplier_pullback_3
 end
 
-"""
-    Divider(x)
-
-Stores a fixed `x` and divides by it, then squares the result.
-
-Especially for testing the gradient of higher order functions with respect to `x`.
-```
-julia> map(Divider(2), [1 2 3 4 10])
-1×5 Matrix{Float64}:
- 0.25  1.0  2.25  4.0  25.0
-```
-"""
-struct Divider{T<:Real}
-    x::T
-end
-(d::Divider)(y::Real) = (y / d.x)^2
-
-function ChainRulesCore.rrule(d::Divider, y::Real)
-    Divider_pullback(dΩ) = (Tangent{typeof(d)}(; x = -2 * dΩ * y^2 / d.x^3), 2 * dΩ * y / d.x^2)
-    return d(y), Divider_pullback
-end
-
 """
     Counter()
 
@@ -110,11 +88,6 @@ end
         test_rrule(Multiplier(1.0 + 2im), 3.0 + 4im, 5.0 - 6im)
         test_rrule(Multiplier(rand(2,3)), rand(3,4), rand(4,5))
     end
-    
-    @testset "Divider" begin
-        test_rrule(Divider(2.3), 4.5)
-        test_rrule(Divider(0.2), -3.4)
-    end
 
     @testset "Counter" begin
         c = Counter()