Merge pull request #489 from mcabbott/reverse

Add rules for `reverse`, `circshift`

Author: mcabbott

src/rulesets/Base/array.jl

@@ -262,19 +262,55 @@ function rrule(::typeof(hvcat), rows, values::Union{AbstractArray, Number}...)
 end
 
 #####
-##### `fill`
+##### `reverse`
 #####
 
-function rrule(::typeof(fill), value::Any, dims::Tuple{Vararg{Int}})
-    function fill_pullback(Ȳ)
-        return (NoTangent(), sum(Ȳ), NoTangent())
+# 1-dim case allows start/stop, N-dim case takes dims keyword
+# whose defaults changed in Julia 1.6... just pass them all through:
+
+function frule((_, xdot), ::typeof(reverse), x::AbstractArray, args...; kw...)
+    return reverse(x, args...; kw...), reverse(xdot, args...; kw...)
+end
+
+function rrule(::typeof(reverse), x::AbstractArray, args...; kw...)
+    project = ProjectTo(x)
+    nots = map(_ -> NoTangent(), args)
+    function reverse_pullback(dy)
+        dx = @thunk project(reverse(unthunk(dy), args...; kw...))
+        return (NoTangent(), dx, nots...)
     end
-    return fill(value, dims), fill_pullback
+    return reverse(x, args...; kw...), reverse_pullback
+end
+
+#####
+##### `circshift`
+#####
+
+function frule((_, xdot), ::typeof(circshift), x::AbstractArray, shifts)
+    return circshift(x, shifts), circshift(xdot, shifts)
 end
 
-function rrule(::typeof(fill), value::Any, dims::Int...)
-    function fill_pullback(Ȳ)
-        return (NoTangent(), sum(Ȳ), ntuple(_->NoTangent(), length(dims))...)
+function rrule(::typeof(circshift), x::AbstractArray, shifts)
+    project = ProjectTo(x)
+    function circshift_pullback(dy)
+        dx = @thunk project(circshift(unthunk(dy), map(-, shifts)))
+        # Note that circshift! is useless for InplaceableThunk, as it overwrites completely
+        return (NoTangent(), dx, NoTangent())
     end
-    return fill(value, dims), fill_pullback
+    return circshift(x, shifts), circshift_pullback
+end
+
+#####
+##### `fill`
+#####
+
+function frule((_, xdot), ::typeof(fill), x::Any, dims...)
+    return fill(x, dims...), fill(xdot, dims...)
+end
+
+function rrule(::typeof(fill), x::Any, dims...)
+    project = x isa Union{Number, AbstractArray{<:Number}} ? ProjectTo(x) : identity
+    nots = map(_ -> NoTangent(), dims)
+    fill_pullback(Ȳ) = (NoTangent(), project(sum(Ȳ)), nots...)
+    return fill(x, dims...), fill_pullback
 end

test/rulesets/Base/array.jl

@@ -101,7 +101,52 @@ end
     test_rrule(hvcat, 1, rand(3)', transpose(rand(3)) ⊢ rand(1,3); check_inferred=VERSION>v"1.1")
 end
 
+@testset "reverse" begin
+    # Forward
+    test_frule(reverse, rand(5))
+    test_frule(reverse, rand(5), 2, 4)
+    test_frule(reverse, rand(5), fkwargs=(dims=1,))
+
+    test_frule(reverse, rand(3,4), fkwargs=(dims=2,))
+    if VERSION >= v"1.6"
+        test_frule(reverse, rand(3,4))
+        test_frule(reverse, rand(3,4,5), fkwargs=(dims=(1,3),))
+    end
+
+    # Reverse
+    test_rrule(reverse, rand(5))
+    test_rrule(reverse, rand(5), 2, 4)
+    test_rrule(reverse, rand(5), fkwargs=(dims=1,))
+
+    test_rrule(reverse, rand(3,4), fkwargs=(dims=2,))
+    if VERSION >= v"1.6"
+        test_rrule(reverse, rand(3,4))
+        test_rrule(reverse, rand(3,4,5), fkwargs=(dims=(1,3),))
+
+        # Structured
+        y, pb = rrule(reverse, Diagonal([1,2,3]))
+        @test unthunk(pb(rand(3,3))[2]) isa Diagonal
+    end
+end
+
+@testset "circshift" begin
+    # Forward
+    test_frule(circshift, rand(10), 1)
+    test_frule(circshift, rand(10), (1,))
+    test_frule(circshift, rand(3,4), (-7,2))
+
+    # Reverse
+    test_rrule(circshift, rand(10), 1)
+    test_rrule(circshift, rand(10) .+ im, -2)
+    test_rrule(circshift, rand(10), (1,))
+    test_rrule(circshift, rand(3,4), (-7,2))
+end
+
 @testset "fill" begin
-    test_rrule(fill, 44.0, 4; check_inferred=false)
-    test_rrule(fill, 2.0, (3, 3, 3))
+    test_frule(fill, 12.3, 4)
+    test_frule(fill, 5.0, (6, 7))
+
+    test_rrule(fill, 44.4, 4)
+    test_rrule(fill, 55 + 0.5im, 5)
+    test_rrule(fill, 3.3, (3, 3, 3))
 end