Add rule for cumprod

[mcabbott]

Closes #254.

The approach here much like that in #335 for prod, and quite different to FluxML/Zygote.jl#294 .

[mcabbott]

Timing these today, there isn't in fact a clear winner, it depends on the size, and how many zeros are encountered. This approach uses less memory. Some numbers (on two computers):

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(10,100)))[1];
  6.625 μs (10 allocations: 16.01 KiB)   # m1 mac + rosetta, Julia 1.6
  20.720 μs (21 allocations: 16.49 KiB)  # xeon

julia> x=rand(10,100); x[1:21:end].=0; # half the columns have a zero

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $x)[1];
  6.883 μs (8 allocations: 15.98 KiB)   # m1
  18.045 μs (21 allocations: 16.49 KiB) # xeon

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(3,100)))[1];  # fewer rows
  1.625 μs (8 allocations: 5.10 KiB)  # m1
  7.330 μs (21 allocations: 5.62 KiB) # xeon

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(30,100)))[1];  # more rows
  36.750 μs (10 allocations: 47.13 KiB)  # m1
  138.382 μs (23 allocations: 47.65 KiB) # xeon

Compare to FluxML/Zygote.jl#294 (which could probably be optimised a bit)

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(10,100)))[1];  # no zeros -- about the same
  6.367 μs (11 allocations: 60.02 KiB)   # m1
  20.141 μs (11 allocations: 60.02 KiB)  # xeon

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $x)[1];  # half with a zero -- much slower
  86.583 μs (774 allocations: 289.86 KiB)
  277.610 μs (774 allocations: 289.86 KiB)
  
julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(3,100)))[1];  # few rows -- slower
  3.531 μs (11 allocations: 21.88 KiB)  # m1

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(30,100)))[1];  # more rows -- faster
  15.333 μs (18 allocations: 169.34 KiB) # m1

The last case, cumprod along a dimension of size 30, gives lots of tiny numbers like 1e-16. My guess is that this is most useful for products of a few numbers.

Edit, after 70334fc, which adds a simpler path for the case of no zeros (like the Zygote PR):

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(10,100)))[1];
  4.863 μs (6 allocations: 23.86 KiB)    # m1
  15.388 μs (19 allocations: 24.38 KiB)  # xeon

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(3,100)))[1];  # fewer rows
  2.000 μs (6 allocations: 7.55 KiB)
  8.765 μs (19 allocations: 8.06 KiB)

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(30,100)))[1];  # more rows
  12.125 μs (9 allocations: 70.59 KiB)
  37.687 μs (22 allocations: 71.11 KiB)

Edit', after b75f94f, the original path is now faster than the fast path was:

julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(10,100)))[1];
  3.036 μs (8 allocations: 15.98 KiB)
  11.527 μs (21 allocations: 16.49 KiB)
  
julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $x)[1]; # half with a zero
  2.866 μs (8 allocations: 15.98 KiB)
  10.990 μs (21 allocations: 16.49 KiB)
   
julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(3,100)))[1];  # fewer rows
  1.471 μs (8 allocations: 5.10 KiB)
  6.963 μs (21 allocations: 5.62 KiB)
  
julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(30,100)))[1];  # more rows
  7.903 μs (10 allocations: 47.13 KiB)
  25.996 μs (23 allocations: 47.65 KiB)

[codecov-commenter]

Codecov Report

Merging #420 (ab8384e) into master (afd4cfb) will decrease coverage by 0.05%.
The diff coverage is 96.29%.

@@            Coverage Diff             @@
##           master     #420      +/-   ##
==========================================
- Coverage   98.51%   98.46%   -0.06%     
==========================================
  Files          21       21              
  Lines        2094     2148      +54     
==========================================
+ Hits         2063     2115      +52     
- Misses         31       33       +2     

Impacted Files	Coverage Δ
src/rulesets/Base/mapreduce.jl	98.05% <96.29%> (-0.95%)	⬇️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update afd4cfb...ab8384e. Read the comment docs.

[oxinabox]

why are we testing values here?
It would be good to add comments explaining what we are particularly checking for that test_rrule will not catch

[mcabbott]

test_rrule doesn't seem to accept integer input, this tests that the rule still does.

julia> test_rrule(cumprod, [1,2,3])
test_rrule: cumprod on Vector{Int64}: Error During Test at /Users/me/.julia/packages/ChainRulesTestUtils/6oOem/src/testers.jl:227
  Got exception outside of a @test
  InexactError: Int64(0.99)

[oxinabox]

add comment to that effect?

Yeah can't test methods that require integers with finite differencing.
Since once you apply a finite difference you get a Float64 instead.
Which means you don't hit the method for integers.

[mcabbott]

The comment now says "# rule allows integer input, but test_rrule does not"

[mcabbott]

Good to go? No rush except that something may change under it...

[oxinabox]

Sorry i lost track of this one.
LGTM