fix mean bug

Author: Miha Zgubic

Project.toml

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

src/rulesets/Statistics/statistics.jl

@@ -2,9 +2,9 @@
 ##### `mean`
 #####
 
+_denom(x, dims) = size(x, dims)
 _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)
+_denom(x, dims::Union{Tuple, AbstractArray}) = 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