Merge pull request #20 from JuliaDiff/npr/fdm0.9

Update FiniteDifferences

Author: nickrobinson251

Project.toml

@@ -1,16 +1,18 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ChainRulesCore = "0.7"
-FiniteDifferences = "0.7, 0.8, 0.9"
+ChainRulesCore = "0.7.1"
+Compat = "3"
+FiniteDifferences = "0.9"
 julia = "1"
 
 [extras]

src/ChainRulesTestUtils.jl

@@ -2,7 +2,7 @@ module ChainRulesTestUtils
 
 using ChainRulesCore
 using ChainRulesCore: frule, rrule
-using ChainRulesCore: AbstractDifferential
+using Compat: only
 using FiniteDifferences
 using LinearAlgebra
 using Test
@@ -11,8 +11,12 @@ const _fdm = central_fdm(5, 1)
 
 export test_scalar, frule_test, rrule_test, isapprox, generate_well_conditioned_matrix
 
+# TODO: reconsider these https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/7
+Base.isapprox(a, b::Union{AbstractZero, AbstractThunk}; kwargs...) = isapprox(b, a; kwargs...)
+Base.isapprox(d_ad::AbstractThunk, d_fd; kwargs...) = isapprox(extern(d_ad), d_fd; kwargs...)
 Base.isapprox(d_ad::DoesNotExist, d_fd; kwargs...) = error("Tried to differentiate w.r.t. a `DoesNotExist`")
-Base.isapprox(d_ad::AbstractDifferential, d_fd; kwargs...) = isapprox(extern(d_ad), d_fd; kwargs...)
+# Call `all` to handle the case where `Zero` is standing in for a non-scalar zero
+Base.isapprox(d_ad::Zero, d_fd; kwargs...) = all(isapprox.(extern(d_ad), d_fd; kwargs...))
 
 """
     _make_fdm_call(fdm, f, ȳ, xs, ignores) -> Tuple
@@ -49,7 +53,6 @@ function _make_fdm_call(fdm, f, ȳ, xs, ignores)
     end
     fdexpr = :(j′vp($fdm, $sig -> $call, $ȳ, $(newxs...)))
     fd = eval(fdexpr)
-    fd isa Tuple || (fd = (fd,))
     args = Any[nothing for _ in 1:length(xs)]
     for (dx, ind) in zip(fd, arginds)
         args[ind] = dx
@@ -170,7 +173,7 @@ function rrule_test(f, ȳ, (x, x̄)::Tuple{Any, Any}; rtol=1e-9, atol=1e-9, fdm
 
     @test ∂self === NO_FIELDS  # No internal fields
     # Correctness testing via finite differencing.
-    x̄_fd = j′vp(fdm, f, ȳ, x)
+    x̄_fd = only(j′vp(fdm, f, ȳ, x))  # j′vp returns a tuple, but `f` is a unary function.
     @test isapprox(x̄_ad, x̄_fd; rtol=rtol, atol=atol, kwargs...)
 end
 

test/runtests.jl

@@ -4,20 +4,48 @@ using Random
 using Test
 
 @testset "ChainRulesTestUtils.jl" begin
-    double(x) = 2x
-    @scalar_rule(double(x), 2)
-    test_scalar(double, 2)
-
-    fst(x, y) = x
-    ChainRulesCore.frule((_, dx, dy), ::typeof(fst), x, y) = (x, dx)
+    @testset "test_scalar" begin
+        double(x) = 2x
+        @scalar_rule(double(x), 2)
+        test_scalar(double, 2)
+    end
 
-    function ChainRulesCore.rrule(::typeof(fst), x, y)
-        function fst_pullback(Δx)
-            return (NO_FIELDS, Δx, Zero())
+    @testset "unary: identity(x)" begin
+        function ChainRulesCore.frule((_, ẏ), ::typeof(identity), x)
+            return x, ẏ
+        end
+        function ChainRulesCore.rrule(::typeof(identity), x)
+            function identity_pullback(ȳ)
+                return (NO_FIELDS, ȳ)
+            end
+            return x, identity_pullback
+        end
+        @testset "frule_test" begin
+            frule_test(identity, (randn(), randn()))
+            frule_test(identity, (randn(4), randn(4)))
+        end
+        @testset "rrule_test" begin
+            rrule_test(identity, randn(), (randn(), randn()))
+            rrule_test(identity, randn(4), (randn(4), randn(4)))
         end
-        return x, fst_pullback
     end
 
-    frule_test(fst, (2, 4.0), (3, 5.0))
-    rrule_test(fst, rand(), (2.0, 4.0), (3.0, 5.0))
+    @testset "binary: fst(x, y)" begin
+        fst(x, y) = x
+        ChainRulesCore.frule((_, dx, dy), ::typeof(fst), x, y) = (x, dx)
+        function ChainRulesCore.rrule(::typeof(fst), x, y)
+            function fst_pullback(Δx)
+                return (NO_FIELDS, Δx, Zero())
+            end
+            return x, fst_pullback
+        end
+        @testset "frule_test" begin
+            frule_test(fst, (2, 4.0), (3, 5.0))
+            frule_test(fst, (randn(4), randn(4)), (randn(4), randn(4)))
+        end
+        @testset "rrule_test" begin
+            rrule_test(fst, rand(), (2.0, 4.0), (3.0, 5.0))
+            rrule_test(fst, randn(4), (randn(4), randn(4)), (randn(4), randn(4)))
+        end
+    end
 end