grad for reverse

FluxML · Mar 26, 2019 · f5edcc5 · f5edcc5
1 parent 38e3a88
commit f5edcc5
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/src/lib/array.jl b/src/lib/array.jl
@@ -240,6 +240,13 @@ Base.permutedims(xs::TrackedArray, perm) = track(permutedims, xs, perm)
 Base.PermutedDimsArray(xs::TrackedArray, perm) = track(PermutedDimsArray, xs, perm)
 @grad PermutedDimsArray(xs, perm) = PermutedDimsArray(data(xs), perm), Δ -> (PermutedDimsArray(Δ, invperm(perm)),nothing)
 
+Base.reverse(xs::TrackedArray; dims) = track(reverse, xs, dims = dims)
+@grad reverse(xs; dims) = reverse(data(xs), dims = dims), Δ -> (reverse(Δ, dims = dims), nothing)
+Base.reverse(xs::TrackedVector) = track(reverse, xs)
+@grad reverse(xs) = reverse(data(xs)), Δ -> (reverse(Δ),)
+Base.reverse(xs::TrackedVector, start, stop) = track(reverse, xs, start, stop)
+@grad reverse(xs, start, stop) = reverse(data(xs), start, stop), Δ -> (reverse(Δ, start, stop), nothing, nothing)
+
 function _kron(mat1::AbstractMatrix,mat2::AbstractMatrix)
     m1, n1 = size(mat1)
     mat1_rsh = reshape(mat1,(1,m1,1,n1))

diff --git a/test/tracker.jl b/test/tracker.jl
@@ -125,6 +125,10 @@ end
 @test gradtest(x -> permutedims(x, [3,1,2]), rand(4,5,6))
 @test gradtest(x -> PermutedDimsArray(x, [3,1,2]), rand(4,5,6))
 
+@test gradtest(reverse, rand(5))
+@test gradtest(x -> reverse(x, dims=2), rand(4,5,6))
+@test gradtest(x -> reverse(x, 2, 4), rand(5))
+
 @test gradtest(x -> repeat(x; inner=2), rand(5))
 @test gradtest(x -> repeat(x; inner=2, outer=3), rand(5))
 @test gradtest(x -> repeat(x; inner=(2,2,1), outer=(1,1,3)), rand(5,4,3))