fixup

Author: mcabbott

src/rulesets/Base/mapreduce.jl

@@ -226,7 +226,8 @@ end
 function rrule(::typeof(cumprod), x::AbstractVector{<:Real}; dims::Integer=1)
     y = cumprod(x; dims=dims)  # does nothing unless dims == 1
     project_x = ProjectTo(x)
-    function cumprod_pullback_1(dy)
+    function cumprod_pullback_1(dy_raw)
+        dy = unthunk(dy_raw)
         dx_thunk = InplaceableThunk(
             dx -> if dims == 1
                 ∇cumprod!(dx, x, dy, y)
@@ -240,15 +241,16 @@ function rrule(::typeof(cumprod), x::AbstractVector{<:Real}; dims::Integer=1)
                 dy
             end)
             )
-        return (NO_FIELDS, dx_thunk)
+        return (NoTangent(), dx_thunk)
     end
     return y, cumprod_pullback_1
 end
 
 function rrule(::typeof(cumprod), x::AbstractArray{<:Real}; dims::Integer)
     y = cumprod(x; dims=dims)
     project_x = ProjectTo(x)
-    function cumprod_pullback_2(dy)
+    function cumprod_pullback_2(dy_raw)
+        dy = unthunk(dy_raw)
         dx_thunk = InplaceableThunk(
             dx -> if dims <= ndims(x)
                 vald = Val(Int(dims))
@@ -264,7 +266,7 @@ function rrule(::typeof(cumprod), x::AbstractArray{<:Real}; dims::Integer)
                 dy
             end)
             )
-        return (NO_FIELDS, dx_thunk)
+        return (NoTangent(), dx_thunk)
     end
     return y, cumprod_pullback_2
 end

test/rulesets/Base/mapreduce.jl

@@ -187,7 +187,7 @@ end
             @test_throws Exception cumprod(Symmetric([1 2; 3 4]), dims=1) # forward pass fails, so can't test gradient
 
             back = rrule(cumprod, Diagonal([1, 2]); dims=1)[2]
-            @test unthunk(back(fill(0.5, 2, 2))[2]) ≈ [1/2 3/2; 1/2 0]
+            @test unthunk(back(fill(0.5, 2, 2))[2]) ≈ [1/2 0; 0 0]  # ProjectTo'd to Diagonal now
         end
     end
 end