Enable testing of rules with @not_implemented (#140)

Co-authored-by: Lyndon White <[email protected]>

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesTestUtils"
 uuid = "cdddcdb0-9152-4a09-a978-84456f9df70a"
-version = "0.6.7"
+version = "0.6.8"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
@@ -11,7 +11,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-ChainRulesCore = "0.9.13"
+ChainRulesCore = "0.9.39"
 Compat = "3"
 FiniteDifferences = "0.12"
 julia = "1"

src/check_result.jl

@@ -33,6 +33,18 @@ check_equal(x::Zero, y::Zero; kwargs...) = @test true
 check_equal(x::DoesNotExist, y::Nothing; kwargs...) = @test true
 check_equal(x::Nothing, y::DoesNotExist; kwargs...) = @test true
 
+# Checking equality with `NotImplemented` reports `@test_broken` since the derivative has intentionally
+# not yet been implemented
+# `@test_broken x == y` yields more descriptive messages than `@test_broken false`
+check_equal(x::ChainRulesCore.NotImplemented, y; kwargs...) = @test_broken x == y
+check_equal(x, y::ChainRulesCore.NotImplemented; kwargs...) = @test_broken x == y
+# In this case we check for equality (messages etc. have to be equal)
+function check_equal(
+    x::ChainRulesCore.NotImplemented, y::ChainRulesCore.NotImplemented; kwargs...
+)
+    return @test x == y
+end
+
 """
     _can_pass_early(actual, expected; kwargs...)
 Used to check if `actual` is basically equal to `expected`, so we don't need to check deeper;
@@ -137,3 +149,20 @@ function _check_add!!_behaviour(acc, val; kwargs...)
     acc_mutated = deepcopy(acc)  # prevent this test changing others
     check_equal(add!!(acc_mutated, val), acc + val; kwargs...)
 end
+
+# Checking equality with `NotImplemented` reports `@test_broken` since the derivative has intentionally
+# not yet been implemented
+# `@test_broken x == y` yields more descriptive messages than `@test_broken false`
+function _check_add!!_behaviour(acc_mutated, acc::ChainRulesCore.NotImplemented; kwargs...)
+    return @test_broken acc_mutated == acc
+end
+function _check_add!!_behaviour(acc_mutated::ChainRulesCore.NotImplemented, acc; kwargs...)
+    return @test_broken acc_mutated == acc
+end
+# In this case we check for equality (messages etc. have to be equal)
+function _check_add!!_behaviour(
+    acc_mutated::ChainRulesCore.NotImplemented, acc::ChainRulesCore.NotImplemented;
+    kwargs...,
+)
+    return @test acc_mutated == acc
+end

test/check_result.jl

@@ -33,7 +33,7 @@ end
 
     @testset "check_equal" begin
 
-        @testset "possive cases" begin
+        @testset "positive cases" begin
             check_equal(1.0, 1.0)
             check_equal(1.0 + im, 1.0 + im)
             check_equal(1.0, 1.0+1e-10)  # isapprox _behaviour
@@ -46,6 +46,10 @@ end
 
             check_equal(@thunk(10*0.1*[[1.0], [2.0]]), [[1.0], [2.0]])
 
+            check_equal(@not_implemented(""), rand(3))
+            check_equal(rand(3), @not_implemented(""))
+            check_equal(@not_implemented("a"), @not_implemented("a"))
+
             check_equal(
                 Composite{Tuple{Float64, Float64}}(1.0, 2.0),
                 Composite{Tuple{Float64, Float64}}(1.0, 2.0)
@@ -92,6 +96,8 @@ end
             @test fails(()->check_equal([[1.0], [2.0]], [[1.1], [2.0]]))
 
             @test fails(()->check_equal(@thunk(10*[[1.0], [2.0]]), [[1.0], [2.0]]))
+
+            @test fails(()->check_equal(@not_implemented("a"), @not_implemented("b")))
         end
         @testset "type negative" begin
             @test fails() do  # these have different primals so should not be equal

test/testers.jl

@@ -14,7 +14,6 @@ f_noninferrable_pullback(x) = x
 f_noninferrable_thunk(x, y) = x + y
 f_inferrable_pullback_only(x) = x > 0 ? Float64(x) : Float32(x)
 
-
 function finplace!(x; y = [1])
     y[1] = 2
     x .*= y[1]
@@ -520,4 +519,11 @@ end
         end
         @test errors(() -> test_rrule(foo, [1.0, 2.0, 3.0], 2), "should use DoesNotExist()")
     end
+
+    @testset "NotImplemented" begin
+        f_notimplemented(x, y) = (x + y, x - y)
+        @scalar_rule f_notimplemented(x, y) (@not_implemented(""), 1) (1, -1)
+        test_frule(f_notimplemented, randn(), randn())
+        test_rrule(f_notimplemented, randn(), randn())
+    end
 end