Merge pull request #628 from JuliaDiff/mz/repeat

Fix corner case for repeat

Author: mzgubic

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.35.1"
+version = "1.35.2"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/array.jl

@@ -166,18 +166,19 @@ function frule((_, ẋs), ::typeof(repeat), xs::AbstractArray, cnt...; kw...)
     return repeat(xs, cnt...; kw...), repeat(ẋs, cnt...; kw...)
 end
 
-function rrule(::typeof(repeat), xs::AbstractArray; inner=ntuple(Returns(1), ndims(xs)), outer=ntuple(Returns(1), ndims(xs)))
+function rrule(::typeof(repeat), xs::AbstractArray; inner=nothing, outer=nothing)
 
     project_Xs = ProjectTo(xs)
     S = size(xs)
+    inner_size = inner === nothing ? ntuple(Returns(1), ndims(xs)) : inner
     function repeat_pullback(ȳ)
         dY = unthunk(ȳ)
         Δ′ = zero(xs)
         # Loop through each element of Δ, calculate source dimensions, accumulate into Δ′
         for (dest_idx, val) in pairs(IndexCartesian(), dY)
-            # First, round dest_idx[dim] to nearest gridpoint defined by inner[dim], then
+            # First, round dest_idx[dim] to nearest gridpoint defined by inner_dims[dim], then
             # wrap around based on original size S.
-            src_idx = [mod1(div(dest_idx[dim] - 1, inner[dim]) + 1, S[dim]) for dim in 1:length(S)]
+            src_idx = [mod1(div(dest_idx[dim] - 1, inner_size[dim]) + 1, S[dim]) for dim in 1:length(S)]
             Δ′[src_idx...] += val
         end
         x̄ = project_Xs(Δ′)

test/rulesets/Base/array.jl

@@ -128,6 +128,7 @@ end
     test_rrule(repeat, rand(4, 5))
     test_rrule(repeat, rand(4, 5); fkwargs = (outer=(1,2),))
     test_rrule(repeat, rand(4, 5); fkwargs = (inner=(1,2), outer=(1,3)))
+    test_rrule(repeat, rand(4, 5); fkwargs = (outer=2,))
 
     test_rrule(repeat, rand(4, ), 2)
     test_rrule(repeat, rand(4, 5), 2)