add implementation

Author: Miha Zgubic

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,40 @@
+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)
+            for (Δy, i) in zip(Δys, inds)
+                Δxs[i] += Δy
+            end
+            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(xs))),
+            sort_add!
+        )
+
+        return NO_FIELDS, Δxs
+    end
+    return ys, sort_pullback
+end