Move mean rule to a Statistics ruleset folder

Author: nickrobinson251

src/ChainRules.jl

@@ -29,6 +29,8 @@ include("rulesets/Base/array.jl")
 include("rulesets/Base/broadcast.jl")
 include("rulesets/Base/mapreduce.jl")
 
+include("rulesets/Statistics/statistics.jl")
+
 include("rulesets/LinearAlgebra/utils.jl")
 include("rulesets/LinearAlgebra/blas.jl")
 include("rulesets/LinearAlgebra/dense.jl")

src/rulesets/Base/mapreduce.jl

@@ -62,25 +62,3 @@ function rrule(::typeof(sum), ::typeof(abs2), x::AbstractArray{<:Real}; dims=:)
     ∂x = Rule(ȳ -> 2ȳ .* x)
     return y, (DNERule(), ∂x)
 end
-
-#####
-##### `mean`
-#####
-
-_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
-
-function rrule(::typeof(mean), x::AbstractArray{<:Real}; dims=:)
-    _, dx = rrule(sum, x; dims=dims)
-    n = _denom(x, dims)
-    return mean(x; dims=dims), Rule(ȳ -> dx(ȳ) / n)
-end
-
-function rrule(::typeof(mean), f, x::AbstractArray{<:Real})
-    _, (_, dx) = rrule(sum, f, x)
-    n = _denom(x, :)
-    return mean(f, x), (DNERule(), Rule(ȳ -> dx(ȳ) / n))
-end

src/rulesets/Statistics/statistics.jl

@@ -0,0 +1,21 @@
+#####
+##### `mean`
+#####
+
+_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
+
+function rrule(::typeof(mean), x::AbstractArray{<:Real}; dims=:)
+    _, dx = rrule(sum, x; dims=dims)
+    n = _denom(x, dims)
+    return mean(x; dims=dims), Rule(ȳ -> dx(ȳ) / n)
+end
+
+function rrule(::typeof(mean), f, x::AbstractArray{<:Real})
+    _, (_, dx) = rrule(sum, f, x)
+    n = _denom(x, :)
+    return mean(f, x), (DNERule(), Rule(ȳ -> dx(ȳ) / n))
+end

test/rulesets/Base/mapreduce.jl

@@ -63,16 +63,5 @@
             x̄_fd = j′vp(central_fdm(5, 1), x->sum(x, dims=2), ȳ, X)
             @test x̄_ad ≈ x̄_fd atol=1e-9 rtol=1e-9
         end
-    end
-    @testset "mean" begin
-        rng = MersenneTwister(999)
-        n = 9
-        rrule_test(mean, randn(rng), (abs2, nothing), (randn(rng, n), randn(rng, n)))
-        X = randn(rng, n, n)
-        y, dX = rrule(mean, X; dims=1)
-        ȳ = randn(rng, size(y))
-        X̄_ad = dX(ȳ)
-        X̄_fd = j′vp(central_fdm(5, 1), x->mean(x, dims=1), ȳ, X)
-        @test X̄_ad ≈ X̄_fd rtol=1e-9 atol=1e-9
-    end
+    end  # sum
 end

test/rulesets/Statistics/statistics.jl

@@ -0,0 +1,11 @@
+@testset "mean" begin
+    rng = MersenneTwister(999)
+    n = 9
+    rrule_test(mean, randn(rng), (abs2, nothing), (randn(rng, n), randn(rng, n)))
+    X = randn(rng, n, n)
+    y, dX = rrule(mean, X; dims=1)
+    ȳ = randn(rng, size(y))
+    X̄_ad = dX(ȳ)
+    X̄_fd = j′vp(central_fdm(5, 1), x->mean(x, dims=1), ȳ, X)
+    @test X̄_ad ≈ X̄_fd rtol=1e-9 atol=1e-9
+end

test/runtests.jl

@@ -27,6 +27,10 @@ include("test_util.jl")
             include(joinpath("rulesets", "Base", "broadcast.jl"))
         end
 
+        @testset "Statistics" begin
+            include(joinpath("rulesets", "Statistics", "statistics.jl"))
+        end
+
         @testset "LinearAlgebra" begin
             include(joinpath("rulesets", "LinearAlgebra", "dense.jl"))
             include(joinpath("rulesets", "LinearAlgebra", "structured.jl"))