code review suggestions

Author: Miha Zgubic

src/rulesets/Base/sort.jl

@@ -4,18 +4,15 @@ function rrule(::typeof(partialsort), xs::AbstractVector, k::Union{Integer,Ordin
 
     function partialsort_pullback(Δys)
         function partialsort_add!(Δxs)
-            for (Δy, i) in zip(Δys, inds)
-                Δxs[i] += Δy
-            end
+            Δxs[inds] += Δys
             return Δxs
         end
 
-        Δxs = InplaceableThunk(
-            @thunk(partialsort_add!(zero(xs))),
-            partialsort_add!
-        )
+        Δxs = InplaceableThunk(@thunk(partialsort_add!(zero(xs))), partialsort_add!)
+
         return NO_FIELDS, Δxs, DoesNotExist()
     end
+
     return ys, partialsort_pullback
 end
 
@@ -29,10 +26,7 @@ function rrule(::typeof(sort), xs::AbstractVector; kwargs...)
             return Δxs
         end
 
-        Δxs = InplaceableThunk(
-            @thunk(sort_add!(zero(xs))),
-            sort_add!
-        )
+        Δxs = InplaceableThunk(@thunk(sort_add!(zero(Δys))), sort_add!)
 
         return NO_FIELDS, Δxs
     end