Rules for var, std (#560)

* rules for var, std

* Apply 2 suggestions

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

Author: mcabbott

src/rulesets/Statistics/statistics.jl

@@ -18,3 +18,39 @@ function rrule(::typeof(mean), x::AbstractArray{<:Real}; dims=:)
     end
     return y_sum / n, mean_pullback
 end
+
+#####
+##### variance
+#####
+
+function rrule(
+    ::typeof(Statistics.var),
+    x::AbstractArray{<:Number};
+    corrected::Bool=true,
+    dims=:,
+    mean=mean(x, dims=dims)
+)
+    y = Statistics.var(x; corrected=corrected, mean=mean, dims=dims)
+    function variance_pullback(dy)
+        pre = 2 // (_denom(x, dims) - corrected)
+        dx = pre .* unthunk(dy) .* (x .- mean)
+        return (NoTangent(), ProjectTo(x)(dx))
+    end
+    y, variance_pullback
+end
+
+function rrule(
+    ::typeof(Statistics.std),
+    x::AbstractArray{<:Number}; 
+    corrected::Bool=true,
+    dims=:,
+    mean=mean(x, dims=dims)
+)
+    y = Statistics.std(x; corrected=corrected, mean=mean, dims=dims)
+    function std_pullback(dy)
+        pre = 1 // (_denom(x, dims) - corrected)
+        dx = pre .* unthunk(dy) .* (x .- mean) ./ y
+        return (NoTangent(), ProjectTo(x)(dx))
+    end
+    y, std_pullback
+end

test/rulesets/Statistics/statistics.jl

@@ -8,3 +8,18 @@
         test_rrule(mean, randn(n,4); fkwargs=(;dims=2))
     end
 end
+
+@testset "variation: $var" for var in (std, var)
+    test_rrule(var, randn(3))
+    test_rrule(var, randn(4, 5); fkwargs=(; corrected=false))
+    test_rrule(var, randn(ComplexF64, 6))
+    test_rrule(var, Diagonal(randn(6)))
+
+    test_rrule(var, randn(4, 5); fkwargs=(; dims=1))
+    test_rrule(var, randn(ComplexF64, 4, 5); fkwargs=(; dims=2, corrected=false))
+    test_rrule(var, UpperTriangular(randn(5, 5)); fkwargs=(; dims=1))
+
+    x = PermutedDimsArray(randn(3, 4, 5), (3, 2, 1))
+    xm = mean(x; dims=(1, 3))
+    test_rrule(var, x; fkwargs=(; dims=(1, 3), mean=xm))
+end