Add rrules for mapreduce, mapfoldr, mapfoldl

Author: ararslan

src/rules/mapreduce.jl

@@ -15,6 +15,30 @@ 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`
 #####

test/rules/mapreduce.jl

@@ -8,4 +8,27 @@
         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
 end