Rule for mean(f,x) (#615)

* mean(f,x) rule

* Revert changes to sum_f_x

* simpler rule, stiffer test

* Apply 2 suggestions

Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com>

Co-authored-by: Miha Zgubic <mzgubic@users.noreply.github.com>

Author: mcabbott

src/rulesets/Statistics/statistics.jl

@@ -6,19 +6,31 @@ _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)
 
-# 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)
+function rrule(::typeof(mean), x::AbstractArray{<:Union{Real,Complex,AbstractArray}}; dims=:)
+    y_sum, sum_pullback = rrule(sum, x; dims)
     n = _denom(x, dims)
     function mean_pullback(ȳ)
-        _, ∂sum_x = sum_pullback(ȳ)
-        ∂x = unthunk(∂sum_x) / n
+        _, ∂x = sum_pullback(unthunk(ȳ) / 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,11 +1,32 @@
 @testset "mean" begin
-    n = 9
-    @testset "Basic" begin
-        test_rrule(mean, randn(n))
+    @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)
     end
     @testset "with dims kwargs" begin
-        test_rrule(mean, randn(n); fkwargs=(;dims=1))
-        test_rrule(mean, randn(n,4); fkwargs=(;dims=2))
+        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,)))
     end
 end
 

test/test_helpers.jl

@@ -35,6 +35,28 @@ 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()
 
@@ -88,6 +110,11 @@ 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()