Merge pull request JuliaDiff#59 from JuliaDiff/aa/mapreduce

Add a few more reduction rrules

Author: ararslan

Project.toml

@@ -7,6 +7,7 @@ Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
 Cassette = "^0.2"

src/ChainRules.jl

@@ -3,6 +3,7 @@ module ChainRules
 using Cassette
 using LinearAlgebra
 using LinearAlgebra.BLAS
+using Statistics
 using Base.Broadcast: materialize, materialize!, broadcasted, Broadcasted, broadcastable
 
 if VERSION < v"1.3.0-DEV.142"

src/rules/base.jl

@@ -77,3 +77,7 @@
 frule(::typeof(*), x, y) = x * y, Rule((Δx, Δy) -> Δx * y + x * Δy)
 
 rrule(::typeof(*), x, y) = x * y, (Rule(ΔΩ -> ΔΩ * y'), Rule(ΔΩ -> x' * ΔΩ))
+
+frule(::typeof(identity), x) = x, Rule(identity)
+
+rrule(::typeof(identity), x) = x, Rule(identity)

src/rules/mapreduce.jl

@@ -15,10 +15,72 @@ function rrule(::typeof(map), f, xs...)
     return y, (DNERule(), ∂xs...)
 end
 
+#####
+##### `mapreduce`, `mapfoldl`, `mapfoldr`
+#####
+
+for mf in (:mapreduce, :mapfoldl, :mapfoldr)
+    sig = :(rrule(::typeof($mf), f, op, x::AbstractArray{<:Real}))
+    call = :($mf(f, op, x))
+    if mf === :mapreduce
+        insert!(sig.args, 2, Expr(:parameters, Expr(:kw, :dims, :(:))))
+        insert!(call.args, 2, Expr(:parameters, Expr(:kw, :dims, :dims)))
+    end
+    body = quote
+        y = $call
+        ∂x = Rule() do ȳ
+            broadcast(x, ȳ) do xi, ȳi
+                _, ∂xi = _checked_rrule(f, xi)
+                extern(∂xi(ȳi))
+            end
+        end
+        return y, (DNERule(), DNERule(), ∂x)
+    end
+    eval(Expr(:function, sig, body))
+end
+
 #####
 ##### `sum`
 #####
 
 frule(::typeof(sum), x) = (sum(x), Rule(sum))
 
 rrule(::typeof(sum), x) = (sum(x), Rule(cast))
+
+function rrule(::typeof(sum), f, x::AbstractArray{<:Real}; dims=:)
+    y, (_, _, ∂x) = rrule(mapreduce, f, Base.add_sum, x; dims=dims)
+    return y, (DNERule(), ∂x)
+end
+
+function rrule(::typeof(sum), x::AbstractArray{<:Real}; dims=:)
+    y, (_, ∂x) = rrule(sum, identity, x; dims=dims)
+    return y, ∂x
+end
+
+function rrule(::typeof(sum), ::typeof(abs2), x::AbstractArray{<:Real}; dims=:)
+    y = sum(abs2, x; dims=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

test/rules/base.jl

@@ -97,5 +97,11 @@ end
             @test dy === x / h * cy.value[2]
         end
     end
+    @testset "identity" begin
+        rng = MersenneTwister(1)
+        n = 4
+        rrule_test(identity, randn(rng), (randn(rng), randn(rng)))
+        rrule_test(identity, randn(rng, 4), (randn(rng, 4), randn(rng, 4)))
+    end
 end
 # TODO: Non-trig stuff

test/rules/linalg/dense.jl

@@ -4,23 +4,6 @@ function generate_well_conditioned_matrix(rng, N)
 end
 
 @testset "linalg" begin
-    @testset "sum" begin
-        @testset "Vector" begin
-            rng, M = MersenneTwister(123456), 3
-            frule_test(sum, (randn(rng, M), randn(rng, M)))
-            rrule_test(sum, randn(rng), (randn(rng, M), randn(rng, M)))
-        end
-        @testset "Matrix" begin
-            rng, M, N = MersenneTwister(123456), 3, 4
-            frule_test(sum, (randn(rng, M, N), randn(rng, M, N)))
-            rrule_test(sum, randn(rng), (randn(rng, M, N), randn(rng, M, N)))
-        end
-        @testset "Array{T, 3}" begin
-            rng, M, N, P = MersenneTwister(123456), 3, 7, 11
-            frule_test(sum, (randn(rng, M, N, P), randn(rng, M, N, P)))
-            rrule_test(sum, randn(rng), (randn(rng, M, N, P), randn(rng, M, N, P)))
-        end
-    end
     @testset "dot" begin
         @testset "Vector" begin
             rng, M = MersenneTwister(123456), 3

test/rules/mapreduce.jl

@@ -8,4 +8,71 @@
         rrule_test(map, ȳ, (sin, nothing), (x, vx))
         rrule_test(map, ȳ, (+, nothing), (x, vx), (randn(rng, n), randn(rng, n)))
     end
+    @testset "mapreduce" begin
+        rng = MersenneTwister(6)
+        n = 10
+        x = randn(rng, n)
+        vx = randn(rng, n)
+        ȳ = randn(rng)
+        rrule_test(mapreduce, ȳ, (sin, nothing), (+, nothing), (x, vx))
+        # With keyword arguments (not yet supported in rrule_test)
+        X = randn(rng, n, n)
+        y, (_, _, dx) = rrule(mapreduce, abs2, +, X; dims=2)
+        ȳ = randn(rng, size(y))
+        x̄_ad = dx(ȳ)
+        x̄_fd = j′vp(central_fdm(5, 1), x->mapreduce(abs2, +, x; dims=2), ȳ, X)
+        @test x̄_ad ≈ x̄_fd atol=1e-9 rtol=1e-9
+    end
+    @testset "$f" for f in (mapfoldl, mapfoldr)
+        rng = MersenneTwister(10)
+        n = 7
+        x = randn(rng, n)
+        vx = randn(rng, n)
+        ȳ = randn(rng)
+        rrule_test(f, ȳ, (cos, nothing), (+, nothing), (x, vx))
+    end
+    @testset "sum" begin
+        @testset "Vector" begin
+            rng, M = MersenneTwister(123456), 3
+            frule_test(sum, (randn(rng, M), randn(rng, M)))
+            rrule_test(sum, randn(rng), (randn(rng, M), randn(rng, M)))
+        end
+        @testset "Matrix" begin
+            rng, M, N = MersenneTwister(123456), 3, 4
+            frule_test(sum, (randn(rng, M, N), randn(rng, M, N)))
+            rrule_test(sum, randn(rng), (randn(rng, M, N), randn(rng, M, N)))
+        end
+        @testset "Array{T, 3}" begin
+            rng, M, N, P = MersenneTwister(123456), 3, 7, 11
+            frule_test(sum, (randn(rng, M, N, P), randn(rng, M, N, P)))
+            rrule_test(sum, randn(rng), (randn(rng, M, N, P), randn(rng, M, N, P)))
+        end
+        @testset "function argument" begin
+            rng = MersenneTwister(1)
+            n = 8
+            rrule_test(sum, randn(rng), (cos, nothing), (randn(rng, n), randn(rng, n)))
+            rrule_test(sum, randn(rng), (abs2, nothing), (randn(rng, n), randn(rng, n)))
+        end
+        @testset "keyword arguments" begin
+            rng = MersenneTwister(33)
+            n = 4
+            X = randn(rng, n, n)
+            y, dX = rrule(sum, X; dims=2)
+            ȳ = randn(rng, size(y))
+            x̄_ad = dX(ȳ)
+            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

test/runtests.jl

@@ -1,6 +1,6 @@
 # TODO: more tests!
 
-using ChainRules, Test, FDM, LinearAlgebra, LinearAlgebra.BLAS, Random
+using ChainRules, Test, FDM, LinearAlgebra, LinearAlgebra.BLAS, Random, Statistics
 using ChainRules: extern, accumulate, accumulate!, store!, @scalar_rule,
     Wirtinger, wirtinger_primal, wirtinger_conjugate, add_wirtinger, mul_wirtinger,
     Zero, add_zero, mul_zero, One, add_one, mul_one, Casted, cast, add_casted, mul_casted,