Support SpecialFunctions 0.8

also improves tests.
Also adds gradient for sign.

Author: oxinabox

test/rulesets/Base/base.jl

@@ -52,9 +52,9 @@
             test_scalar(acotd, 1/x)
         end
         @testset "Multivariate" begin
-            x, y = rand(2)
             @testset "atan2" begin
                 # https://en.wikipedia.org/wiki/Atan2
+                x, y = rand(2)
                 ratan = atan(x, y)
                 u = x^2 + y^2
                 datan = y/u - 2x/u
@@ -71,19 +71,11 @@
             end
 
             @testset "sincos" begin
-                rsincos = sincos(x)
-                dsincos = cos(x) - 2sin(x)
-
-                r, pushforward = frule(sincos, x)
-                @test r === rsincos
-                df1, df2 = pushforward(NamedTuple(), 1)
-                @test df1 + 2df2 === dsincos
-
-                r, pullback = rrule(sincos, x)
-                @test r === rsincos
-                ds, df = pullback(1, 2)
-                @test df === dsincos
-                @test ds === NO_FIELDS
+                x, Δx, x̄ = randn(3)
+                Δz = (randn(), randn())
+
+                frule_test(sincos, (x, Δx))
+                rrule_test(sincos, Δz, (x, x̄))
             end
         end
     end  # Trig
@@ -114,17 +106,16 @@
     end
 
     @testset "Unary complex functions" begin
-        for x in (-6, rand.((Float32, Float64, Complex{Float32}, Complex{Float64}))...)
-            rtol = x isa Complex{Float32} ? 1e-6 : 1e-9
-            test_scalar(real, x; rtol=rtol)
-            test_scalar(imag, x; rtol=rtol)
+        for x in (-4.1, 6.4, 1.0+0.5im, -10.0+1.5im)
+            test_scalar(real, x)
+            test_scalar(imag, x)
 
-            test_scalar(abs, x; rtol=rtol)
-            test_scalar(hypot, x; rtol=rtol)
+            test_scalar(abs, x)
+            test_scalar(hypot, x)
 
-            test_scalar(angle, x; rtol=rtol)
-            test_scalar(abs2, x; rtol=rtol)
-            test_scalar(conj, x; rtol=rtol)
+            test_scalar(angle, x)
+            test_scalar(abs2, x)
+            test_scalar(conj, x)
         end
     end
 
@@ -146,14 +137,14 @@
         test_accumulation(rand(2, 5), dy)
     end
 
-    @testset "hypot(x, y)" begin
+    @testset "binary trig ($f)" for f in (hypot, atan)
         rng = MersenneTwister(123456)
-        x, Δx, x̄ = randn(rng, 3)
+        x, Δx, x̄ = 10randn(rng, 3)
         y, Δy, ȳ = randn(rng, 3)
         Δz = randn(rng)
 
-        frule_test(hypot, (x, Δx), (y, Δy))
-        rrule_test(hypot, Δz, (x, x̄), (y, ȳ))
+        frule_test(f, (x, Δx), (y, Δy))
+        rrule_test(f, Δz, (x, x̄), (y, ȳ))
     end
 
     @testset "identity" begin
@@ -175,7 +166,7 @@
         @testset "Zero over the point discontinuity" begin
             # Can't do finite differencing because we are lying
             # following the subgradient convention.
-            
+
             _, pb = rrule(sign, 0.0)
             _, x̄ = pb(10.5)
             @test extern(x̄) == 0

test/rulesets/packages/SpecialFunctions.jl

@@ -1,5 +1,5 @@
 using SpecialFunctions
-#=
+
 @testset "SpecialFunctions" for x in (1, -1, 0, 0.5, 10, -17.1, 1.5 + 0.7im)
     test_scalar(SpecialFunctions.erf, x)
     test_scalar(SpecialFunctions.erfc, x)
@@ -33,7 +33,7 @@ using SpecialFunctions
         test_scalar(SpecialFunctions.trigamma, x)
     end
 end
-=#
+
 @testset "log gamma and co" begin
     # SpecialFunctions 0.7->0.8 changes:
     for x in (1.5, 2.5, 10.5, -0.6, -2.6, 1.6+1.6im, 1.6-1.6im, -4.6+1.6im)
@@ -47,7 +47,12 @@ end
 
         if isdefined(SpecialFunctions, :logabsgamma)
             isreal(x) || continue
-            
+
+            x, Δx, x̄ = randn(3)
+            Δz = (randn(), randn())
+
+            frule_test(SpecialFunctions.logabsgamma, (x, Δx))
+            rrule_test(SpecialFunctions.logabsgamma, Δz, (x, x̄))
         end
     end
 

test/runtests.jl

@@ -20,13 +20,14 @@ println("Testing ChainRules.jl")
 @testset "ChainRules" begin
     include("helper_functions.jl")
     @testset "rulesets" begin
+
         @testset "Base" begin
             include(joinpath("rulesets", "Base", "base.jl"))
-            #include(joinpath("rulesets", "Base", "array.jl"))
-            #include(joinpath("rulesets", "Base", "mapreduce.jl"))
-            #include(joinpath("rulesets", "Base", "broadcast.jl"))
+            include(joinpath("rulesets", "Base", "array.jl"))
+            include(joinpath("rulesets", "Base", "mapreduce.jl"))
+            include(joinpath("rulesets", "Base", "broadcast.jl"))
         end
-        #==
+
         print(" ")
 
         @testset "Statistics" begin
@@ -41,7 +42,6 @@ println("Testing ChainRules.jl")
             include(joinpath("rulesets", "LinearAlgebra", "factorization.jl"))
             include(joinpath("rulesets", "LinearAlgebra", "blas.jl"))
         end
-        ==#
 
         print(" ")
 

test/test_util.jl

@@ -98,7 +98,13 @@ function frule_test(f, xẋs::Tuple{Any, Any}...; rtol=1e-9, atol=1e-9, fdm=_fdm
 
     # Correctness testing via finite differencing.
     dΩ_fd = jvp(fdm, xs->f(xs...), (xs, ẋs))
-    @test isapprox(dΩ_ad, dΩ_fd; rtol=rtol, atol=atol, kwargs...)
+    @test isapprox(
+        collect(dΩ_ad),  # Use collect so can use vector equality
+        collect(dΩ_fd);
+        rtol=rtol,
+        atol=atol,
+        kwargs...
+    )
 end
 
 
@@ -108,6 +114,7 @@ end
 # Arguments
 - `f`: Function to which rule should be applied.
 - `ȳ`: adjoint w.r.t. output of `f` (should generally be set randomly).
+  Should be same structure as `f(x)` (so if multiple returns should be a tuple)
 - `x`: input at which to evaluate `f` (should generally be set to an arbitary point in the domain).
 - `x̄`: currently accumulated adjoint (should generally be set randomly).
 
@@ -118,8 +125,15 @@ function rrule_test(f, ȳ, (x, x̄)::Tuple{Any, Any}; rtol=1e-9, atol=1e-9, fdm
 
     # Check correctness of evaluation.
     fx, pullback = ChainRules.rrule(f, x)
-    @test fx ≈ f(x)
-    (∂self, x̄_ad) = pullback(ȳ)
+    @test collect(fx) ≈ collect(f(x))  # use collect so can do vector equality
+    (∂self, x̄_ad) = if fx isa Tuple
+        # If the function returned multiple values,
+        # then it must have multiple seeds for propagating backwards
+        pullback(ȳ...)
+    else
+        pullback(ȳ)
+    end
+
     @test ∂self === NO_FIELDS  # No internal fields
     # Correctness testing via finite differencing.
     x̄_fd = j′vp(fdm, f, ȳ, x)