Update src/rulesets/Base/mapreduce.jl

Co-authored-by: Michael Abbott <[email protected]>

Author: Keno

src/rulesets/Base/mapreduce.jl

@@ -7,12 +7,12 @@ function frule((_, ẋ), ::typeof(sum), x; dims=:)
 end
 
 function sum_rrule(x, dims::Colon)
-    y = sum(x; dims=dims)
-    let xdims=size(x)
+    y = sum(x)
+    let x_axes = axes(x)
         function sum_pullback(ȳ)
             x̄ = InplaceableThunk(
                 x -> x .+= ȳ,
-                @thunk(fill(ȳ, xdims...)),
+                @thunk(fill(ȳ, x_axes)),
             )
             return (NoTangent(), x̄)
         end