bang

Author: mcabbott

src/rulesets/Base/mapreduce.jl

@@ -162,11 +162,16 @@ function frule((_, xdot), ::typeof(cumsum), x::AbstractArray; dims::Integer)
 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)
+        if ChainRulesCore.is_inplaceable_destination(step1) && VERSION >= v"1.6"
             step2 = cumsum!(step1, step1; dims=dims)
             step3 = reverse!(step2; dims=dims)
         else

test/rulesets/Base/mapreduce.jl

@@ -240,11 +240,13 @@ end
 @testset "Accumulations" begin
     @testset "cumsum" begin
         v = round.(10 .* randn(9), sigdigits=3)
-        m = round.(10 .* randn(4,5), sigdigits=3)
+        m = round.(10 .* randn(4, 5), sigdigits=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)