sum which allows 2nd derivatives

Author: mcabbott

src/rulesets/Base/mapreduce.jl

@@ -1,24 +1,48 @@
 #####
-##### `sum`
+##### `sum(x)`
 #####
 
 function frule((_, ẋ), ::typeof(sum), x; dims=:)
     return sum(x; dims=dims), sum(ẋ; dims=dims)
 end
 
-function rrule(::typeof(sum), x::AbstractArray{T}; dims=:) where {T<:Number}
+function rrule(::typeof(sum), x::AbstractArray; dims=:)
+    project = ProjectTo(x)
     y = sum(x; dims=dims)
-    function sum_pullback(ȳ)
-        # broadcasting the two works out the size no-matter `dims`
-        x̄ = InplaceableThunk(
-            x -> x .+= ȳ,
-            @thunk(broadcast(last∘tuple, x, ȳ)),
+    function sum_pullback(dy_raw)
+        dy = unthunk(dy_raw)
+        x_thunk = InplaceableThunk(
+            # Protect `dy` from broadcasting, for when `x` is an array of arrays:
+            dx -> dx .+= (dims isa Colon ? Ref(dy) : dy),
+            @thunk project(_unsum(x, dy, dims))  # `_unsum` handles Ref internally
         )
-        return (NoTangent(), x̄)
+        return (NoTangent(), x_thunk)
     end
     return y, sum_pullback
 end
 
+# This broadcasts `dy` to the shape of `x`, and should preserve e.g. CuArrays, StaticArrays.
+# Ideally this would only need `typeof(x)` not `x`, but `similar` only has a suitable method
+# when `eltype(x) == eltype(dy)`, which isn't guaranteed.
+_unsum(x, dy, dims) = broadcast(last∘tuple, x, dy)
+_unsum(x, dy, ::Colon) = broadcast(last∘tuple, x, Ref(dy))
+
+# Allow for second derivatives of `sum`, by writing rules for `_unsum`:
+
+function frule((_, _, dydot, _), ::typeof(_unsum), x, dy, dims)
+    return _unsum(x, dy, dims), _unsum(x, dydot, dims)
+end
+
+function rrule(::typeof(_unsum), x, dy, dims)
+    z = _unsum(x, dy, dims)
+    _unsum_pullback(dz) = (NoTangent(), NoTangent(), sum(unthunk(dz); dims=dims), NoTangent())
+    return z, _unsum_pullback
+end
+
+#####
+##### `sum(f, x)`
+#####
+
 # Can't map over Adjoint/Transpose Vector
 function rrule(
     config::RuleConfig{>:HasReverseMode},

test/rulesets/Base/mapreduce.jl

@@ -1,12 +1,32 @@
 @testset "Maps and Reductions" begin
-    @testset "sum" begin
-        sizes = (3, 4, 7)
-        @testset "dims = $dims" for dims in (:, 1)
-            @testset "Array{$N, $T}" for N in eachindex(sizes), T in (Float64, ComplexF64)
-                x = randn(T, sizes[1:N]...)
-                test_frule(sum, x; fkwargs=(;dims=dims))
-                test_rrule(sum, x; fkwargs=(;dims=dims))
-            end
+    @testset "sum(x; dims=$dims)" for dims in (:, 2, (1,3))
+        # Forward
+        test_frule(sum, rand(5); fkwargs=(;dims=dims))
+        test_frule(sum, rand(ComplexF64, 2,3,4); fkwargs=(;dims=dims))
+
+        # Reverse
+        test_rrule(sum, rand(5); fkwargs=(;dims=dims))
+        test_rrule(sum, rand(ComplexF64, 2,3,4); fkwargs=(;dims=dims))
+
+        # Structured matrices
+        test_rrule(sum, rand(5)'; fkwargs=(;dims=dims))
+        y, back = rrule(sum, UpperTriangular(rand(5,5)); dims=dims)
+        unthunk(back(y*(1+im))[2]) isa UpperTriangular{Float64}
+
+        # Function allowing for 2nd derivatives
+        for x in (rand(5), rand(2,3,4))
+            dy = maximum(x; dims=dims)
+            test_frule(ChainRules._unsum, x, dy, dims)
+            test_rrule(ChainRules._unsum, x, dy, dims)
+        end
+
+        # Arrays of arrays
+        for x  in ([rand(ComplexF64, 3) for _ in 1:4], [rand(3) for _ in 1:2, _ in 1:3, _ in 1:4])
+            test_rrule(sum, x; fkwargs=(;dims=dims), check_inferred=false)
+
+            dy = sum(x; dims=dims)
+            ddy = rrule(ChainRules._unsum, x, dy, dims)[2](x)[3]
+            @test size(ddy) == size(dy)
         end
     end