partialsort! rule tests

EnzymeAD · Mar 31, 2024 · 82da102 · 82da102
1 parent 38c63fe
commit 82da102
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 3 deletions.
diff --git a/test/internal_rules.jl b/test/internal_rules.jl
@@ -44,6 +44,28 @@ end
     @test autodiff(Forward, f2, BatchDuplicated(2.0, (1.0, 2.0)))[1] == (var"1"=-3.0, var"2"=-6.0)
     @test autodiff(Reverse, f2, Active, Active(2.0))[1][1] == -3
 
+    function f3(x)
+        a = [2.0, 2.5, x, 1.0]
+        return partialsort(a, 2)
+    end
+
+    @test autodiff(Forward, f3, Duplicated(1.5, 1.0))[1] == 1.0
+    @test autodiff(Forward, f3, BatchDuplicated(1.5, (1.0, 2.0)))[1] == (var"1"=1.0, var"2"=2.0)
+    @test autodiff(Reverse, f3, Active(1.5))[1][1] == 1.0
+    @test autodiff(Reverse, f3, Active(2.5))[1][1] == 0.0
+
+    function f4(x)
+        a = [2.0, 2.5, x, x / 2]
+        y = partialsort(a, 1:2)
+        return sum(y)
+    end
+
+    @test autodiff(Forward, f4, Duplicated(1.5, 1.0))[1] == 1.5
+    @test autodiff(Forward, f4, BatchDuplicated(1.5, (1.0, 2.0)))[1] == (var"1"=1.5, var"2"=3.0)
+    @test autodiff(Reverse, f4, Active(1.5))[1][1] == 1.5
+    @test autodiff(Reverse, f4, Active(4.0))[1][1] == 0.5
+    @test autodiff(Reverse, f4, Active(6.0))[1][1] == 0.0
+
     dd = Duplicated([TPair(1, 2), TPair(2, 3), TPair(0, 1)], [TPair(0, 0), TPair(0, 0), TPair(0, 0)])
     res = Enzyme.autodiff(Reverse, sorterrfn, dd, Active(1.0))
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -2882,11 +2882,10 @@ end
     @test autodiff(Forward, f6, Duplicated(4.0, 1.0))[1]   ≈ 5/3
 
     f7(x) = median([2.0, 1.0, x])
-    # Fails on Julia 1.9 due to #880
-    #=@test autodiff(Reverse, f7, Active, Active(1.5))[1][1] == 1
+    @test autodiff(Reverse, f7, Active, Active(1.5))[1][1] == 1
     @test autodiff(Forward, f7, Duplicated(1.5, 1.0))[1]   == 1
     @test autodiff(Reverse, f7, Active, Active(2.5))[1][1] == 0
-    @test autodiff(Forward, f7, Duplicated(2.5, 1.0))[1]   == 0=#
+    @test autodiff(Forward, f7, Duplicated(2.5, 1.0))[1]   == 0
 
     f8(x) = middle([2.0, x, 1.0])
     @test autodiff(Reverse, f8, Active, Active(2.5))[1][1] == 0.5