Merge branch 'main' into mz/nograd

Author: Miha Zgubic

src/rulesets/Base/indexing.jl

@@ -93,6 +93,7 @@ function ∇eachslice(dys_raw, x::AbstractArray, vd::Val{dim}) where {dim}
     end
     return ProjectTo(x)(dx)
 end
+∇eachslice(dys::AbstractZero, x::AbstractArray, vd::Val{dim}) where {dim} = dys
 
 _zero_fill!(dx::AbstractArray{<:Number}) = fill!(dx, zero(eltype(dx)))
 _zero_fill!(dx::AbstractArray) = map!(zero, dx, dx)

src/rulesets/Base/mapreduce.jl

@@ -186,25 +186,28 @@ end
 #####
 
 function frule((_, xdot), ::typeof(cumsum), x::AbstractArray; dims::Integer)
-    return cumsum(x; dims=dims), cumsum(xdot; dims=dims)
+    return cumsum(x; dims), cumsum(xdot; 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)
+    return cumsum!(y, x; dims), cumsum!(ydot, xdot; 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)
+function rrule(::typeof(cumsum), x::AbstractArray{T,N}; dims::Integer) where {T,N}
     project = ProjectTo(x)
     function cumsum_pullback(dy)
+        if dims > N  # trivial case, for which reverse fails
+            return (NoTangent(), project(unthunk(dy)))
+        end
         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)
+        if ChainRulesCore.is_inplaceable_destination(step1)
+            step2 = cumsum!(step1, step1; dims)
+            step3 = reverse!(step2; dims)
         else
-            step2 = cumsum(step1; dims=dims)
-            step3 = reverse(step2; dims=dims)
+            step2 = cumsum(step1; dims)
+            step3 = reverse(step2; dims)
         end
         return (NoTangent(), project(step3))
     end

test/rulesets/Base/mapreduce.jl

@@ -277,6 +277,7 @@ end
         test_rrule(cumsum, v)
         test_rrule(cumsum, v; fkwargs=(;dims=1))
         test_rrule(cumsum, m; fkwargs=(;dims=2))
+        test_rrule(cumsum, m; fkwargs=(;dims=3))  # trivial
     end
     @testset "cumprod" begin
         v = round.(10 .* randn(9), sigdigits=3)