Allow N-dimensional arrays in sorting rules

JuliaDiff · Aug 30, 2023 · a9e831b · a9e831b
1 parent 79722bf
commit a9e831b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/src/rulesets/Base/sort.jl b/src/rulesets/Base/sort.jl
@@ -25,12 +25,12 @@ function rrule(::typeof(partialsort), xs::AbstractVector, k::Union{Integer,Ordin
     return ys, partialsort_pullback
 end
 
-function frule((_, ẋs), ::typeof(sort), xs::AbstractVector; kw...)
+function frule((_, ẋs), ::typeof(sort), xs::AbstractArray; kw...)
     inds = sortperm(xs; kw...)
     return xs[inds], ẋs[inds]
 end
 
-function rrule(::typeof(sort), xs::AbstractVector; kwargs...)
+function rrule(::typeof(sort), xs::AbstractArray; kwargs...)
     inds = sortperm(xs; kwargs...)
     ys = xs[inds]
 

diff --git a/test/rulesets/Base/sort.jl b/test/rulesets/Base/sort.jl
@@ -7,6 +7,15 @@
         # rev
         test_rrule(sort, a)
         test_rrule(sort, a; fkwargs=(;rev=true))
+
+        a = rand(5, 4)
+        for dims in (1, 2)
+            # fwd
+            test_frule(sort, a; fkwargs=(;dims))
+            test_frule(sort, a; fkwargs=(;dims, rev=true))
+            # rev
+            test_rrule(sort, a; fkwargs=(;dims))
+            test_rrule(sort, a; fkwargs=(;dims, rev=true))
     end
     @testset "partialsort" begin
         a = rand(10)