Rules for cumsum (#573)

* cumsum, with variants, and mystery test failures

* fixup

* bang

* digits

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

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

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.20.1" 
+version = "1.21" 
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/mapreduce.jl

@@ -153,6 +153,37 @@ for Config in (RuleConfig, RuleConfig{>:HasReverseMode})
     end
 end
 
+#####
+##### `cumsum`
+#####
+
+function frule((_, xdot), ::typeof(cumsum), x::AbstractArray; dims::Integer)
+    return cumsum(x; dims=dims), cumsum(xdot; dims=dims)
+end
+frule(tang, ::typeof(cumsum), x::AbstractVector) = frule(tang, cumsum, x; dims=1)
+
+function frule((_, ydot, xdot), ::typeof(cumsum!), y::AbstractArray, x::AbstractArray; dims::Integer)
+    return cumsum!(y, x; dims=dims), cumsum!(ydot, xdot; dims=dims)
+end
+frule(t, ::typeof(cumsum!), y::AbstractVector, x::AbstractVector) = frule(t, cumsum!, y, x; dims=1)
+
+function rrule(::typeof(cumsum), x::AbstractArray; dims::Integer)
+    project = ProjectTo(x)
+    function cumsum_pullback(dy)
+        step1 = reverse(unthunk(dy); dims=dims)
+        if ChainRulesCore.is_inplaceable_destination(step1) && VERSION >= v"1.6"
+            step2 = cumsum!(step1, step1; dims=dims)
+            step3 = reverse!(step2; dims=dims)
+        else
+            step2 = cumsum(step1; dims=dims)
+            step3 = reverse(step2; dims=dims)
+        end
+        return (NoTangent(), project(step3))
+    end
+    return cumsum(x; dims=dims), cumsum_pullback
+end
+rrule(::typeof(cumsum), x::AbstractVector) = rrule(cumsum, x; dims=1)
+
 #####
 ##### `prod`
 #####

test/rulesets/Base/mapreduce.jl

@@ -238,6 +238,21 @@ const CFG = ChainRulesTestUtils.ADviaRuleConfig()
 end
 
 @testset "Accumulations" begin
+    @testset "cumsum" begin
+        v = round.(10 .* randn(9), digits=3)
+        m = round.(10 .* randn(4, 5), digits=3)
+
+        # Forward
+        test_frule(cumsum, v)
+        test_frule(cumsum, m; fkwargs=(;dims=1))
+        test_frule(cumsum!, rand(9), v)
+        test_frule(cumsum!, rand(4, 5), m; fkwargs=(;dims=1))
+
+        # Reverse
+        test_rrule(cumsum, v)
+        test_rrule(cumsum, v; fkwargs=(;dims=1))
+        test_rrule(cumsum, m; fkwargs=(;dims=2))
+    end
     @testset "cumprod" begin
         v = round.(10 .* randn(9), sigdigits=3)
         test_rrule(cumprod, v)