From 4ca7fa187589848afb6bb65961e42b4991d4127d Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Thu, 10 Oct 2024 20:29:49 +0200 Subject: [PATCH] Add function filter to type stability checks --- .../src/test_differentiation.jl | 9 +- .../src/tests/type_stability_eval.jl | 359 ++++++++++-------- 2 files changed, 207 insertions(+), 161 deletions(-) diff --git a/DifferentiationInterfaceTest/src/test_differentiation.jl b/DifferentiationInterfaceTest/src/test_differentiation.jl index a10b0b68d..6929033c6 100644 --- a/DifferentiationInterfaceTest/src/test_differentiation.jl +++ b/DifferentiationInterfaceTest/src/test_differentiation.jl @@ -48,6 +48,7 @@ For `type_stability` and `benchmark`, the possible values are `:none`, `:prepare **Type stability options:** - `ignored_modules=nothing`: list of modules that JET.jl should ignore +- `function_filter`: filter for functions that JET.jl should ignore (with a reasonable default) **Benchmark options:** @@ -72,6 +73,11 @@ function test_differentiation( sparsity::Bool=true, # type stability options ignored_modules=nothing, + function_filter=if VERSION >= v"1.11" + @nospecialize(f) -> true + else + @nospecialize(f) -> f != Base.mapreduce_empty # fix for `mapreduce` in jacobian and hessian + end, # benchmark options count_calls::Bool=true, ) @@ -136,7 +142,8 @@ function test_differentiation( adapted_backend, scen; subset=type_stability, - ignored_modules=ignored_modules, + ignored_modules, + function_filter, ) end yield() diff --git a/DifferentiationInterfaceTest/src/tests/type_stability_eval.jl b/DifferentiationInterfaceTest/src/tests/type_stability_eval.jl index 1afb4a2f8..01e173951 100644 --- a/DifferentiationInterfaceTest/src/tests/type_stability_eval.jl +++ b/DifferentiationInterfaceTest/src/tests/type_stability_eval.jl @@ -18,280 +18,319 @@ for op in ALL_OPS if op in [:derivative, :gradient, :jacobian] @eval function test_jet( - ba::AbstractADType, scen::$S1out; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S1out; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, contexts) = deepcopy(scen) prep = $prep_op(f, ba, x, contexts...) (subset == :full) && - @test_opt ignored_modules = ignored_modules $prep_op(f, ba, x, contexts...) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, ba, x, contexts...) (subset == :full) && - @test_opt ignored_modules = ignored_modules $op(f, ba, x, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op( - f, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op(f, prep, ba, x, contexts...) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, prep, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op( - f, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, prep, ba, x, contexts...) return nothing end @eval function test_jet( - ba::AbstractADType, scen::$S1in; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S1in; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, res1, contexts) = deepcopy(scen) prep = $prep_op(f, ba, x, contexts...) (subset == :full) && - @test_opt ignored_modules = ignored_modules $prep_op(f, ba, x, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $op!( - f, res1, ba, x, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op!( - f, res1, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, res1, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, res1, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op!( - f, res1, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, res1, prep, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op!( - f, res1, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, res1, prep, ba, x, contexts...) return nothing end op == :gradient && continue @eval function test_jet( - ba::AbstractADType, scen::$S2out; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S2out; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, y, contexts) = deepcopy(scen) prep = $prep_op(f, y, ba, x, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $prep_op( - f, y, ba, x, contexts... - ) - (subset == :full) && - @test_opt ignored_modules = ignored_modules $op(f, y, ba, x, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op( - f, y, ba, x, contexts... - ) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, y, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, y, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, y, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op( - f, y, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, y, prep, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op( - f, y, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, y, prep, ba, x, contexts...) return nothing end @eval function test_jet( - ba::AbstractADType, scen::$S2in; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S2in; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, y, res1, contexts) = deepcopy(scen) prep = $prep_op(f, y, ba, x, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $prep_op( - f, y, ba, x, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $op!( - f, y, res1, ba, x, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op!( - f, y, res1, ba, x, contexts... - ) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, y, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, y, res1, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, y, res1, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op!( - f, y, res1, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, y, res1, prep, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op!( - f, y, res1, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, y, res1, prep, ba, x, contexts...) return nothing end elseif op in [:second_derivative, :hessian] @eval function test_jet( - ba::AbstractADType, scen::$S1out; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S1out; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, contexts) = deepcopy(scen) prep = $prep_op(f, ba, x, contexts...) (subset == :full) && - @test_opt ignored_modules = ignored_modules $prep_op(f, ba, x, contexts...) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, ba, x, contexts...) (subset == :full) && - @test_opt ignored_modules = ignored_modules $op(f, ba, x, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op( - f, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op(f, prep, ba, x, contexts...) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, prep, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op( - f, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, prep, ba, x, contexts...) return nothing end @eval function test_jet( - ba::AbstractADType, scen::$S1in; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S1in; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, res1, res2, contexts) = deepcopy(scen) prep = $prep_op(f, ba, x, contexts...) (subset == :full) && - @test_opt ignored_modules = ignored_modules $prep_op(f, ba, x, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $op!( - f, res2, ba, x, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op!( - f, res1, res2, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, res2, ba, x, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, res1, res2, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op!( - f, res2, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, res2, prep, ba, x, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op!( - f, res1, res2, prep, ba, x, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, res1, res2, prep, ba, x, contexts...) return nothing end elseif op in [:pushforward, :pullback] @eval function test_jet( - ba::AbstractADType, scen::$S1out; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S1out; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, tang, contexts) = deepcopy(scen) prep = $prep_op(f, ba, x, tang, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $prep_op( - f, ba, x, tang, contexts... - ) - (subset == :full) && - @test_opt ignored_modules = ignored_modules $op(f, ba, x, tang, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op( - f, ba, x, tang, contexts... - ) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op( - f, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, prep, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op( - f, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, prep, ba, x, tang, contexts...) return nothing end @eval function test_jet( - ba::AbstractADType, scen::$S1in; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S1in; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, tang, res1, res2, contexts) = deepcopy(scen) prep = $prep_op(f, ba, x, tang, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $prep_op( - f, ba, x, tang, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $op!( - f, res1, ba, x, tang, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op!( - f, res1, ba, x, tang, contexts... - ) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, res1, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, res1, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op!( - f, res1, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, res1, prep, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op!( - f, res1, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, res1, prep, ba, x, tang, contexts...) return nothing end @eval function test_jet( - ba::AbstractADType, scen::$S2out; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S2out; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, y, tang, contexts) = deepcopy(scen) prep = $prep_op(f, y, ba, x, tang, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $prep_op( - f, y, ba, x, tang, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $op( - f, y, ba, x, tang, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op( - f, y, ba, x, tang, contexts... - ) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, y, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, y, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, y, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op( - f, y, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, y, prep, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op( - f, y, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op(f, y, prep, ba, x, tang, contexts...) return nothing end @eval function test_jet( - ba::AbstractADType, scen::$S2in; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S2in; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, y, tang, res1, contexts) = deepcopy(scen) prep = $prep_op(f, y, ba, x, tang, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $prep_op( - f, y, ba, x, tang, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $op!( - f, y, res1, ba, x, tang, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $val_and_op!( - f, y, res1, ba, x, tang, contexts... - ) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, y, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, y, res1, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, y, res1, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op!( - f, y, res1, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, y, res1, prep, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $val_and_op!( - f, y, res1, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $val_and_op!(f, y, res1, prep, ba, x, tang, contexts...) return nothing end elseif op in [:hvp] @eval function test_jet( - ba::AbstractADType, scen::$S1out; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S1out; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, tang, contexts) = deepcopy(scen) prep = $prep_op(f, ba, x, tang, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $prep_op( - f, ba, x, tang, contexts... - ) (subset == :full) && - @test_opt ignored_modules = ignored_modules $op(f, ba, x, tang, contexts...) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op( - f, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op(f, prep, ba, x, tang, contexts...) return nothing end @eval function test_jet( - ba::AbstractADType, scen::$S1in; subset::Symbol, ignored_modules + ba::AbstractADType, + scen::$S1in; + subset::Symbol, + ignored_modules, + function_filter, ) (; f, x, tang, res1, res2, contexts) = deepcopy(scen) prep = $prep_op(f, ba, x, tang, contexts...) - (subset == :full) && @test_opt ignored_modules = ignored_modules $prep_op( - f, ba, x, tang, contexts... - ) - (subset == :full) && @test_opt ignored_modules = ignored_modules $op!( - f, res2, ba, x, tang, contexts... - ) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $prep_op(f, ba, x, tang, contexts...) + (subset == :full) && + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, res2, ba, x, tang, contexts...) (subset in (:prepared, :full)) && - @test_opt ignored_modules = ignored_modules $op!( - f, res2, prep, ba, x, tang, contexts... - ) + @test_opt ignored_modules = ignored_modules function_filter = + function_filter $op!(f, res2, prep, ba, x, tang, contexts...) return nothing end end