tidy & fix tests

Author: mcabbott

src/rulesets/Base/mapreduce.jl

@@ -247,16 +247,17 @@ end
 function rrule(::typeof(cumprod), x::AbstractArray{<:Real}; dims)
     y = cumprod(x; dims=dims)
     @assert dims isa Integer
-    vald = Val(Int(dims))  # else ∇cumprod_dim! will be type unstable
     function cumprod_pullback_2(dy)
         dx_thunk = InplaceableThunk(
             @thunk if dims <= ndims(x)
+                vald = Val(Int(dims))
                 ∇cumprod_dim(vald, x, dy, y)
             else
                 dy
             end
             ,
             dx -> if dims <= ndims(x)
+                vald = Val(Int(dims))
                 ∇cumprod_dim!(dx, vald, x, dy, y)
             else
                 dx .+= dy
@@ -282,19 +283,6 @@ function ∇cumprod_dim!(dx::AbstractArray, ::Val{dim}, x::AbstractArray, dy, y)
     return dx
 end
 
-#=
-
-julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(10,100)))[1]
-  86.333 μs (2007 allocations: 67.54 KiB)
-
-julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(10,100)))[1]  # with 1 hard-coded
-  5.417 μs (6 allocations: 15.95 KiB)
-
-julia> @btime gradient(x -> sum(cumprod(x, dims=1)), $(rand(10,100)))[1]  # with Val(dim)
-  5.423 μs (6 allocations: 15.95 KiB)
-
-=#
-
 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

@@ -31,7 +31,6 @@
             dy = sum(x; dims=dims)
             ddy = rrule(ChainRules._unsum, x, dy, dims)[2](x)[3]
             @test size(ddy) == size(dy)
-        end
     end
 
     @testset "sum abs2" begin
@@ -155,37 +154,37 @@
     end # prod
 end
 
-
-
+@testset "Accumulations" begin
     @testset "cumprod" begin
-        v = randn(9)
+        v = round.(10 .* randn(9), sigdigits=3)
         test_rrule(cumprod, v)
         v[3] = 0
         test_rrule(cumprod, v)
         v[6] = 0
         test_rrule(cumprod, v)
 
         @testset "higher dimensions, dims=$dims" for dims in (1,2,3)
-            m = rand(4,5)
+            m = round.(10 .* randn(4,5), sigdigits=3)
             test_rrule(cumprod, m; fkwargs=(;dims=dims))
             m[2,2] = 0
             m[2,4] = 0
             test_rrule(cumprod, m; fkwargs=(;dims=dims))
 
-            t = randn(3,3,3)
-            test_rrule(cumprod, x; fkwargs=(;dims=dims))
+            t = round.(10 .* randn(3,3,3), sigdigits=3)
+            test_rrule(cumprod, t; fkwargs=(;dims=dims))
         end
 
         @testset "types" begin
-            back = unthunk(rrule(cumprod, [1, 2, 3])[2])
+            back = unthunk(rrule(cumprod, [1, 2, 3])[2])  # allow integer input
             @test unthunk(back(fill(0.5, 3))[2]) == [9/2, 2, 1]
 
             back = unthunk(rrule(cumprod, PermutedDimsArray([1 2; 3 4], (2,1)); dims=1)[2])
             @test unthunk(back(ones(Float32, 2,2))[2]) == [3 5; 1 3]
 
-            @test_throws Exception cumprod(Symmetric([1 2; 3 4]), dims=1)
+            @test_throws Exception cumprod(Symmetric([1 2; 3 4]), dims=1) # forward pass fails
 
             back = unthunk(rrule(cumprod, Diagonal([1, 2]); dims=1)[2])
             @test unthunk(back(fill(0.5, 2, 2))[2]) ≈ [1/2 3/2; 1/2 0]
         end
     end
+end