Improve call count tests (#82)

* Improve call count tests

* Remove noise

ext/DifferentiationInterfaceChairmarksExt/DifferentiationInterfaceChairmarksExt.jl

@@ -15,14 +15,6 @@ using DifferentiationInterface.DifferentiationTest
 import DifferentiationInterface.DifferentiationTest as DT
 using Test
 
-function soft_test_zero(v)
-    if iszero(v)
-        @test v == 0
-    else
-        @test_broken v == 0
-    end
-end
-
 function DT.run_benchmark(
     backends::Vector{<:AbstractADType},
     operators::Vector{Symbol},
@@ -118,8 +110,8 @@ function benchmark_pushforward_allocating!(
     bench1 = @be zero(dy) value_and_pushforward!(_, ba, f, x, dx, extras)
     bench2 = @be zero(dy) pushforward!(_, ba, f, x, dx, extras)
     if allocations && dy isa Number
-        soft_test_zero(minimum(bench1).allocs)
-        soft_test_zero(minimum(bench2).allocs)
+        @test 0 == minimum(bench1).allocs
+        @test 0 == minimum(bench2).allocs
     end
     record!(data, ba, scen, :value_and_pushforward!, bench1)
     record!(data, ba, scen, :pushforward!, bench2)
@@ -137,7 +129,7 @@ function benchmark_pushforward_mutating!(
         _[1], _[2], ba, f!, x, dx, extras
     )
     if allocations
-        soft_test_zero(minimum(bench1).allocs)
+        @test 0 == minimum(bench1).allocs
     end
     record!(data, ba, scen, :value_and_pushforward!, bench1)
     return nothing
@@ -154,8 +146,8 @@ function benchmark_pullback_allocating!(
     bench1 = @be zero(dx) value_and_pullback!(_, ba, f, x, dy, extras)
     bench2 = @be zero(dx) pullback!(_, ba, f, x, dy, extras)
     if allocations && dy isa Number
-        soft_test_zero(minimum(bench1).allocs)
-        soft_test_zero(minimum(bench2).allocs)
+        @test 0 == minimum(bench1).allocs
+        @test 0 == minimum(bench2).allocs
     end
     record!(data, ba, scen, :value_and_pullback!, bench1)
     record!(data, ba, scen, :pullback!, bench2)
@@ -171,7 +163,7 @@ function benchmark_pullback_mutating!(
     extras = prepare_pullback(ba, f!, x, y)
     bench1 = @be (zero(y), zero(dx)) value_and_pullback!(_[1], _[2], ba, f!, x, dy, extras)
     if allocations
-        soft_test_zero(minimum(bench1).allocs)
+        @test 0 == minimum(bench1).allocs
     end
     record!(data, ba, scen, :value_and_pullback!, bench1)
     return nothing
@@ -186,7 +178,7 @@ function benchmark_derivative_allocating!(
     extras = prepare_derivative(ba, f, x)
     bench1 = @be value_and_derivative(ba, f, x, extras)
     if allocations
-        soft_test_zero(minimum(bench1).allocs)
+        @test 0 == minimum(bench1).allocs
     end
     record!(data, ba, scen, :value_and_derivative, bench1)
     return nothing
@@ -215,7 +207,7 @@ function benchmark_multiderivative_mutating!(
         _[1], _[2], ba, f!, x, extras
     )
     if allocations
-        soft_test_zero(minimum(bench1).allocs)
+        @test 0 == minimum(bench1).allocs
     end
     record!(data, ba, scen, :value_and_multiderivative!, bench1)
     return nothing
@@ -231,8 +223,8 @@ function benchmark_gradient_allocating!(
     bench1 = @be zero(dx) value_and_gradient!(_, ba, f, x, extras)
     bench2 = @be zero(dx) gradient!(_, ba, f, x, extras)
     if allocations
-        soft_test_zero(minimum(bench1).allocs)
-        soft_test_zero(minimum(bench2).allocs)
+        @test 0 == minimum(bench1).allocs
+        @test 0 == minimum(bench2).allocs
     end
     record!(data, ba, scen, :value_and_gradient!, bench1)
     record!(data, ba, scen, :gradient!, bench2)
@@ -264,7 +256,7 @@ function benchmark_jacobian_mutating!(
         _[1], _[2], ba, f!, x, extras
     )
     if allocations
-        soft_test_zero(minimum(bench1).allocs)
+        @test 0 == minimum(bench1).allocs
     end
     record!(data, ba, scen, :value_and_jacobian!, bench1)
     return nothing
@@ -279,7 +271,7 @@ function benchmark_second_derivative_allocating!(
     extras = prepare_second_derivative(ba, f, x)
     bench1 = @be value_derivative_and_second_derivative(ba, f, x, extras)
     if allocations
-        soft_test_zero(minimum(bench1).allocs)
+        @test 0 == minimum(bench1).allocs
     end
     record!(data, ba, scen, :value_derivative_and_second_derivative, bench1)
     return nothing
@@ -297,8 +289,8 @@ function benchmark_hessian_vector_product_allocating!(
     )
     bench2 = @be zero(dx) hessian_vector_product!(_, ba, f, x, dx, extras)
     if allocations  # TODO: distinguish
-        soft_test_zero(minimum(bench1).allocs)
-        soft_test_zero(minimum(bench2).allocs)
+        @test 0 == minimum(bench1).allocs
+        @test 0 == minimum(bench2).allocs
     end
     record!(data, ba, scen, :gradient_and_hessian_vector_product!, bench1)
     record!(data, ba, scen, :hessian_vector_product!, bench2)
@@ -318,8 +310,8 @@ function benchmark_hessian_allocating!(
     )
     bench2 = @be (zero(hess_template)) hessian!(_, ba, f, x, extras)
     if allocations
-        soft_test_zero(minimum(bench1).allocs)
-        soft_test_zero(minimum(bench2).allocs)
+        @test 0 == minimum(bench1).allocs
+        @test 0 == minimum(bench2).allocs
     end
     record!(data, ba, scen, :value_gradient_and_hessian!, bench1)
     record!(data, ba, scen, :hessian!, bench2)

ext/DifferentiationInterfaceJETExt/DifferentiationInterfaceJETExt.jl

@@ -12,7 +12,7 @@ using DifferentiationInterface:
     outer
 using DifferentiationInterface.DifferentiationTest
 import DifferentiationInterface.DifferentiationTest as DT
-using JET: @test_opt
+using JET: @test_call, @test_opt
 using LinearAlgebra: LinearAlgebra
 using Test
 
@@ -100,9 +100,17 @@ function test_type_pushforward_allocating(ba::AbstractADType, scen::Scenario)
     isa(mode(ba), ReverseMode) && return nothing
     (; f, x, dx, dy) = deepcopy(scen)
     dy_in = zero(dy)
+
+    @test_call value_and_pushforward!(dy_in, ba, f, x, dx)
     @test_opt value_and_pushforward!(dy_in, ba, f, x, dx)
+
+    @test_call pushforward!(dy_in, ba, f, x, dx)
     @test_opt pushforward!(dy_in, ba, f, x, dx)
+
+    @test_call value_and_pushforward(ba, f, x, dx)
     @test_opt value_and_pushforward(ba, f, x, dx)
+
+    @test_call pushforward(ba, f, x, dx)
     @test_opt pushforward(ba, f, x, dx)
 end
 
@@ -113,6 +121,8 @@ function test_type_pushforward_mutating(ba::AbstractADType, scen::Scenario)
     f! = f
     y_in = zero(y)
     dy_in = zero(dy)
+
+    @test_call value_and_pushforward!(y_in, dy_in, ba, f!, x, dx)
     @test_opt value_and_pushforward!(y_in, dy_in, ba, f!, x, dx)
 end
 
@@ -122,9 +132,17 @@ function test_type_pullback_allocating(ba::AbstractADType, scen::Scenario)
     isa(mode(ba), ForwardMode) && return nothing
     (; f, x, dx, dy) = deepcopy(scen)
     dx_in = zero(dx)
+
+    @test_call value_and_pullback!(dx_in, ba, f, x, dy)
     @test_opt value_and_pullback!(dx_in, ba, f, x, dy)
+
+    @test_call pullback!(dx_in, ba, f, x, dy)
     @test_opt pullback!(dx_in, ba, f, x, dy)
+
+    @test_call value_and_pullback(ba, f, x, dy)
     @test_opt value_and_pullback(ba, f, x, dy)
+
+    @test_call pullback(ba, f, x, dy)
     @test_opt pullback(ba, f, x, dy)
 end
 
@@ -135,14 +153,20 @@ function test_type_pullback_mutating(ba::AbstractADType, scen::Scenario)
     f! = f
     y_in = zero(y)
     dx_in = zero(dx)
+
+    @test_call value_and_pullback!(y_in, dx_in, ba, f!, x, dy)
     @test_opt value_and_pullback!(y_in, dx_in, ba, f!, x, dy)
 end
 
 ## Derivative
 
 function test_type_derivative_allocating(ba::AbstractADType, scen::Scenario)
     (; f, x) = deepcopy(scen)
+
+    @test_call value_and_derivative(ba, f, x)
     @test_opt value_and_derivative(ba, f, x)
+
+    @test_call derivative(ba, f, x)
     @test_opt derivative(ba, f, x)
 end
 
@@ -151,9 +175,17 @@ end
 function test_type_multiderivative_allocating(ba::AbstractADType, scen::Scenario)
     (; f, x, dy) = deepcopy(scen)
     multider_in = zero(dy)
+
+    @test_call value_and_multiderivative!(multider_in, ba, f, x)
     @test_opt value_and_multiderivative!(multider_in, ba, f, x)
+
+    @test_call multiderivative!(multider_in, ba, f, x)
     @test_opt multiderivative!(multider_in, ba, f, x)
+
+    @test_call value_and_multiderivative(ba, f, x)
     @test_opt value_and_multiderivative(ba, f, x)
+
+    @test_call multiderivative(ba, f, x)
     @test_opt multiderivative(ba, f, x)
 end
 
@@ -163,6 +195,8 @@ function test_type_multiderivative_mutating(ba::AbstractADType, scen::Scenario)
     f! = f
     y_in = zero(y)
     multider_in = zero(dy)
+
+    @test_call value_and_multiderivative!(y_in, multider_in, ba, f!, x)
     @test_opt value_and_multiderivative!(y_in, multider_in, ba, f!, x)
 end
 
@@ -171,9 +205,17 @@ end
 function test_type_gradient_allocating(ba::AbstractADType, scen::Scenario)
     (; f, x, dx) = deepcopy(scen)
     grad_in = zero(dx)
+
+    @test_call value_and_gradient!(grad_in, ba, f, x)
     @test_opt value_and_gradient!(grad_in, ba, f, x)
+
+    @test_call gradient!(grad_in, ba, f, x)
     @test_opt gradient!(grad_in, ba, f, x)
+
+    @test_call value_and_gradient(ba, f, x)
     @test_opt value_and_gradient(ba, f, x)
+
+    @test_call gradient(ba, f, x)
     @test_opt gradient(ba, f, x)
 end
 
@@ -182,9 +224,17 @@ end
 function test_type_jacobian_allocating(ba::AbstractADType, scen::Scenario)
     (; f, x, y) = deepcopy(scen)
     jac_in = zeros(eltype(y), length(y), length(x))
+
+    @test_call value_and_jacobian!(jac_in, ba, f, x)
     @test_opt value_and_jacobian!(jac_in, ba, f, x)
+
+    @test_call jacobian!(jac_in, ba, f, x)
     @test_opt jacobian!(jac_in, ba, f, x)
+
+    @test_call value_and_jacobian(ba, f, x)
     @test_opt value_and_jacobian(ba, f, x)
+
+    @test_call jacobian(ba, f, x)
     @test_opt jacobian(ba, f, x)
 end
 
@@ -194,14 +244,20 @@ function test_type_jacobian_mutating(ba::AbstractADType, scen::Scenario)
     f! = f
     y_in = zero(y)
     jac_in = zeros(eltype(y), length(y), length(x))
+
+    @test_call value_and_jacobian!(y_in, jac_in, ba, f!, x)
     @test_opt value_and_jacobian!(y_in, jac_in, ba, f!, x)
 end
 
 ## Second derivative
 
 function test_type_second_derivative_allocating(ba::AbstractADType, scen::Scenario)
     (; f, x) = deepcopy(scen)
+
+    @test_call value_derivative_and_second_derivative(ba, f, x)
     @test_opt value_derivative_and_second_derivative(ba, f, x)
+
+    @test_call second_derivative(ba, f, x)
     @test_opt second_derivative(ba, f, x)
 end
 
@@ -211,13 +267,27 @@ function test_type_hessian_vector_product_allocating(ba::AbstractADType, scen::S
     (; f, x, dx) = deepcopy(scen)
     grad_in = zero(dx)
     hvp_in = zero(dx)
+
+    @test_call ignored_modules = (LinearAlgebra,) hessian_vector_product!(
+        hvp_in, ba, f, x, dx
+    )
     @test_opt ignored_modules = (LinearAlgebra,) hessian_vector_product!(
         hvp_in, ba, f, x, dx
     )
+
+    @test_call ignored_modules = (LinearAlgebra,) hessian_vector_product(ba, f, x, dx)
     @test_opt ignored_modules = (LinearAlgebra,) hessian_vector_product(ba, f, x, dx)
+
+    @test_call ignored_modules = (LinearAlgebra,) gradient_and_hessian_vector_product!(
+        grad_in, hvp_in, ba, f, x, dx
+    )
     @test_opt ignored_modules = (LinearAlgebra,) gradient_and_hessian_vector_product!(
         grad_in, hvp_in, ba, f, x, dx
     )
+
+    @test_call ignored_modules = (LinearAlgebra,) gradient_and_hessian_vector_product(
+        ba, f, x, dx
+    )
     @test_opt ignored_modules = (LinearAlgebra,) gradient_and_hessian_vector_product(
         ba, f, x, dx
     )
@@ -229,11 +299,21 @@ function test_type_hessian_allocating(ba::AbstractADType, scen::Scenario)
     (; f, x, dx) = deepcopy(scen)
     grad_in = zero(dx)
     hess_in = zeros(eltype(x), length(x), length(x))
+
+    @test_call ignored_modules = (LinearAlgebra,) value_gradient_and_hessian!(
+        grad_in, hess_in, ba, f, x
+    )
     @test_opt ignored_modules = (LinearAlgebra,) value_gradient_and_hessian!(
         grad_in, hess_in, ba, f, x
     )
+
+    @test_call ignored_modules = (LinearAlgebra,) hessian!(hess_in, ba, f, x)
     @test_opt ignored_modules = (LinearAlgebra,) hessian!(hess_in, ba, f, x)
+
+    @test_call ignored_modules = (LinearAlgebra,) value_gradient_and_hessian(ba, f, x)
     @test_opt ignored_modules = (LinearAlgebra,) value_gradient_and_hessian(ba, f, x)
+
+    @test_call ignored_modules = (LinearAlgebra,) hessian(ba, f, x)
     @test_opt ignored_modules = (LinearAlgebra,) hessian(ba, f, x)
 end
 

src/DifferentiationInterface.jl

@@ -29,7 +29,6 @@ include("prepare.jl")
 
 include("pushforward.jl")
 include("pullback.jl")
-include("zero.jl")
 
 include("derivative.jl")
 include("multiderivative.jl")

src/DifferentiationTest/DifferentiationTest.jl

@@ -6,33 +6,35 @@ Testing utilities for [`DifferentiationInterface`](@ref).
 module DifferentiationTest
 
 using ..DifferentiationInterface
+import ..DifferentiationInterface as DI
 using ..DifferentiationInterface:
-    AutoZeroForward,
-    AutoZeroReverse,
     ForwardMode,
+    MutationNotSupported,
     ReverseMode,
     SymbolicMode,
-    MutationNotSupported,
+    inner,
     mode,
     mutation_behavior,
-    inner,
-    outer
+    outer,
+    zero!
 using ADTypes
 using ADTypes: AbstractADType
 using DocStringExtensions
 using Test: @testset, @test
 
-include("utils.jl")
 include("scenario.jl")
 include("benchmark.jl")
 include("call_count.jl")
 include("default_scenarios.jl")
 include("test_operators.jl")
+include("zero.jl")
+include("printing.jl")
 
 export backend_string
 export Scenario, default_scenarios
 export allocating, mutating, scalar_scalar, scalar_array, array_scalar, array_array
 export BenchmarkData, record!
 export test_operators
+export AutoZeroForward, AutoZeroReverse
 
 end