Merge pull request #534 from JuliaDiff/ox/sumbr

Handle sum(f,xs) where f returns a Vector

Author: oxinabox

src/rulesets/Base/mapreduce.jl

@@ -74,8 +74,10 @@ function rrule(
     project = ProjectTo(xs)
 
     function sum_pullback(ȳ)
-        call(f, x) = f(x)  # we need to broadcast this to handle dims kwarg
-        f̄_and_x̄s = call.(pullbacks, ȳ)
+        call(f, x) = f(x)
+        # if dims is :, then need only left-handed only broadcast
+        broadcast_ȳ = dims isa Colon  ? (ȳ,) : ȳ
+        f̄_and_x̄s = call.(pullbacks, broadcast_ȳ)
         # no point thunking as most of work is in f̄_and_x̄s which we need to compute for both
         f̄ = if fieldcount(typeof(f)) === 0 # Then don't need to worry about derivative wrt f
             NoTangent()

test/rulesets/Base/mapreduce.jl

@@ -92,6 +92,15 @@ struct SumRuleConfig <: RuleConfig{Union{HasReverseMode}} end
         test_rrule(sum, sqrt, randn(5,5) .> 0; fkwargs=(;dims=1))
         # ... and Bool produced by function
         @test_skip test_rrule(sum, iszero, randn(5))  # DimensionMismatch("second dimension of A, 1, does not match length of x, 0")
+
+
+        # Functions that return a Vector
+        # see https://github.com/FluxML/Zygote.jl/issues/1074
+        test_rrule(sum, make_two_vec, [1.0, 3.0, 5.0, 7.0])
+        test_rrule(sum, make_two_vec, [1.0 2.0; 3.0 4.0])
+        test_rrule(sum, make_two_vec, [1.0 2.0; 3.0 4.0]; fkwargs=(;dims=2))
+        test_rrule(sum, make_two_vec, [1.0 2.0; 3.0 4.0]; fkwargs=(;dims=1))
+        test_rrule(sum, make_two_vec, [1.0 2.0; 3.0 4.0]; fkwargs=(;dims=(3, 4)))
     end
 
     # https://github.com/JuliaDiff/ChainRules.jl/issues/522

test/test_helpers.jl

@@ -11,8 +11,18 @@ end
 # NoRules - has no rules defined
 struct NoRules; end
 
+"A function that outputs a vector from a scalar for testing"
+make_two_vec(x) = [x, x]
+function ChainRulesCore.rrule(::typeof(make_two_vec), x)
+    make_two_vec_pullback(ȳ) = (NoTangent(), sum(ȳ))
+    return make_two_vec(x), make_two_vec_pullback
+end
+
 @testset "test_helpers.jl" begin
-    @testset "Multiplier functor test-helper" begin
+    @testset "Multiplier functor" begin
         test_rrule(Multiplier(4.0), 3.0)
     end
-end
+    @testset "make_two_vec" begin
+        test_rrule(make_two_vec, 1.5)
+    end
+end