Merge pull request #395 from JuliaDiff/mz/sort

rrule for sort and partialsort

Author: mzgubic

Project.toml

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

src/ChainRules.jl

@@ -36,6 +36,7 @@ include("rulesets/Base/evalpoly.jl")
 include("rulesets/Base/array.jl")
 include("rulesets/Base/arraymath.jl")
 include("rulesets/Base/indexing.jl")
+include("rulesets/Base/sort.jl")
 include("rulesets/Base/mapreduce.jl")
 
 include("rulesets/Statistics/statistics.jl")

src/rulesets/Base/sort.jl

@@ -0,0 +1,34 @@
+function rrule(::typeof(partialsort), xs::AbstractVector, k::Union{Integer,OrdinalRange}; kwargs...)
+    inds = partialsortperm(xs, k; kwargs...)
+    ys = xs[inds]
+
+    function partialsort_pullback(Δys)
+        function partialsort_add!(Δxs)
+            Δxs[inds] += Δys
+            return Δxs
+        end
+
+        Δxs = InplaceableThunk(@thunk(partialsort_add!(zero(xs))), partialsort_add!)
+
+        return NO_FIELDS, Δxs, DoesNotExist()
+    end
+
+    return ys, partialsort_pullback
+end
+
+function rrule(::typeof(sort), xs::AbstractVector; kwargs...)
+    inds = sortperm(xs; kwargs...)
+    ys = xs[inds]
+
+    function sort_pullback(Δys)
+        function sort_add!(Δxs)
+            Δxs[inds] += Δys
+            return Δxs
+        end
+
+        Δxs = InplaceableThunk(@thunk(sort_add!(zero(Δys))), sort_add!)
+
+        return NO_FIELDS, Δxs
+    end
+    return ys, sort_pullback
+end

test/rulesets/Base/sort.jl

@@ -0,0 +1,15 @@
+@testset "sort.jl" begin
+    @testset "sort" begin
+        a = rand(10)
+        test_rrule(sort, a)
+        test_rrule(sort, a; fkwargs=(;rev=true))
+    end
+    @testset "partialsort" begin
+        a = rand(10)
+        test_rrule(partialsort, a, 4 ⊢ nothing)
+        test_rrule(partialsort, a, 3:5 ⊢ nothing)
+        test_rrule(partialsort, a, 1:2:6 ⊢ nothing)
+
+        test_rrule(partialsort, a, 4 ⊢ nothing, fkwargs=(;rev=true))
+    end
+end

test/runtests.jl

@@ -32,6 +32,7 @@ println("Testing ChainRules.jl")
             include_test("rulesets/Base/arraymath.jl")
             include_test("rulesets/Base/indexing.jl")
             include_test("rulesets/Base/mapreduce.jl")
+            include_test("rulesets/Base/sort.jl")
         end
         println()