Check frule primal with isapprox

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/testers.jl

@@ -142,8 +142,12 @@ All keyword arguments except for `fdm` and `fkwargs` are passed to `isapprox`.
 function frule_test(f, xẋs::Tuple{Any, Any}...; rtol=1e-9, atol=1e-9, fdm=_fdm, fkwargs=NamedTuple(), kwargs...)
     _ensure_not_running_on_functor(f, "frule_test")
     xs, ẋs = first.(xẋs), last.(xẋs)
-    Ω, dΩ_ad = frule((NO_FIELDS, ẋs...), f, xs...; fkwargs...)
-    @test f(xs...; fkwargs...) == Ω
+    Ω_ad, dΩ_ad = frule((NO_FIELDS, ẋs...), f, xs...; fkwargs...)
+    Ω = f(xs...; fkwargs...)
+    # if equality check fails, check approximate equality
+    # use collect so can do vector equality
+    # TODO: add isapprox replacement that works for more types
+    @test Ω_ad == Ω || isapprox(collect(Ω_ad), collect(Ω); rtol=rtol, atol=atol)
 
     # Correctness testing via finite differencing.
     dΩ_fd = jvp(fdm, xs->f(xs...; fkwargs...), (xs, ẋs))
@@ -177,8 +181,10 @@ function rrule_test(f, ȳ, xx̄s::Tuple{Any, Any}...; rtol=1e-9, atol=1e-9, fdm
     xs, x̄s = collect(zip(xx̄s...))
     y_ad, pullback = rrule(f, xs...; fkwargs...)
     y = f(xs...; fkwargs...)
+    # if equality check fails, check approximate equality
     # use collect so can do vector equality
-    @test isapprox(collect(y_ad), collect(y); rtol=rtol, atol=atol)
+    # TODO: add isapprox replacement that works for more types
+    @test y_ad == y || isapprox(collect(y_ad), collect(y); rtol=rtol, atol=atol)
     @assert !(isa(ȳ, Thunk))
 
     ∂s = pullback(ȳ)

test/testers.jl

@@ -5,6 +5,8 @@ fbtestkws(x, y; err = true) = err ? error() : x
 
 sinconj(x) = sin(x)
 
+primalapprox(x) = x
+
 @testset "testers.jl" begin
     @testset "test_scalar" begin
         double(x) = 2x
@@ -189,4 +191,18 @@ sinconj(x) = sin(x)
             rrule_test(fbtestkws, randn(4), (randn(4), randn(4)), (randn(4), randn(4)); fkwargs=(; err = false))
         end
     end
+
+    @testset "primal can be only approximately equal" begin
+        ChainRulesCore.frule((_, Δx), ::typeof(primalapprox), x) = (x + sqrt(eps(x)), Δx)
+
+        function ChainRulesCore.rrule(::typeof(primalapprox), x)
+            function primalapprox_pullback(Δx)
+                return (NO_FIELDS, Δx)
+            end
+            return x + sqrt(eps(x)), primalapprox_pullback
+        end
+
+        frule_test(primalapprox, (randn(), randn()); atol = 1e-6)
+        rrule_test(primalapprox, randn(), (randn(), randn()); atol = 1e-6)
+    end
 end