borrow fast path from Zygote 294

Author: mcabbott

src/rulesets/Base/mapreduce.jl

@@ -276,13 +276,28 @@ function ∇cumprod_dim(vald::Val{dim}, x::AbstractArray, dy=fill!(zero(x),1), y
  end
 
 @inline function ∇cumprod_dim!(dx::AbstractArray, ::Val{dim}, x::AbstractArray, dy, y) where {dim}
-    iters = ntuple(k -> k==dim ? Ref(:) : axes(x,k), ndims(x))
-    for ind in Iterators.product(iters...)
-        @views ∇cumprod!(dx[ind...], x[ind...], dy[ind...], y[ind...])
+    if any(iszero, x)
+        iters = ntuple(k -> k==dim ? Ref(:) : axes(x,k), ndims(x))
+        for ind in Iterators.product(iters...)
+            @views ∇cumprod!(dx[ind...], x[ind...], dy[ind...], y[ind...])
+        end
+    else
+        step1 = y .* dy # _rscale!!(y, dy)  # is it safe to mutate y?
+        step2 = _reverse!!(_cumsum!!(_reverse!!(step1, dim), dim), dim)
+        dx .+= step2 ./ x
     end
     return dx
 end
 
+# _rscale!!(A, β) = A .* β
+# _rscale!!(A::StridedArray, β) = A .*= β
+
+_reverse!!(x, dims=1) = reverse(x; dims=dims)
+_reverse!!(x::StridedArray, dims=1) = reverse!(x; dims=dims)
+
+_cumsum!!(x, dims=1) = cumsum(x; dims=dims)
+_cumsum!!(x::StridedArray, dims=1) = cumsum!(x, x; dims=dims)
+
 function ∇cumprod(x::AbstractVector, dy=one(x), y=cumprod(x))
     T = promote_type(eltype(x), eltype(dy))
     dx = fill!(similar(x, T, axes(x)), zero(T))  # axes(x) makes MArray on StaticArrays, Array for structured matrices

test/rulesets/Base/mapreduce.jl

@@ -165,13 +165,16 @@ end
 
         @testset "higher dimensions, dims=$dims" for dims in (1,2,3)
             m = round.(10 .* randn(4,5), sigdigits=3)
-            test_rrule(cumprod, m; fkwargs=(;dims=dims))
+            test_rrule(cumprod, m; fkwargs=(;dims=dims), atol=0.1)
             m[2,2] = 0
             m[2,4] = 0
             test_rrule(cumprod, m; fkwargs=(;dims=dims))
 
             t = round.(10 .* randn(3,3,3), sigdigits=3)
             test_rrule(cumprod, t; fkwargs=(;dims=dims))
+            t[2,2,2] = 0
+            t[2,3,3] = 0
+            test_rrule(cumprod, t; fkwargs=(;dims=dims))
         end
 
         @testset "types" begin