ci: fix tests and update JET harness

Author: MilesCranmerBot

test/Project.toml

@@ -5,6 +5,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
+DynamicExpressions = "a40a106e-89c9-4ca8-8020-a735e8728b6b"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
@@ -26,6 +27,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 Aqua = "0.7"
+JET = "0.11"
 
 [preferences.DynamicExpressions]
 dispatch_doctor_codegen_level = "min"

test/runtests.jl

@@ -23,26 +23,27 @@ if "jet" in test_name
         )
         using JET
         using DynamicExpressions
-        struct MyIgnoredModule
+
+        # JET v0.11 uses the ReportMatcher + match_report API.
+        struct MyIgnoredModule <: JET.ReportMatcher
             mod::Module
         end
-        function JET.match_module(
-            mod::MyIgnoredModule, @nospecialize(report::JET.InferenceErrorReport)
+        function JET.match_report(
+            m::MyIgnoredModule, @nospecialize(report::JET.InferenceErrorReport)
         )
-            s_mod = string(mod.mod)
+            s_mod = string(m.mod)
             any(report.vst) do vst
                 occursin(s_mod, string(JET.linfomod(vst.linfo)))
             end
         end
+
         JET.test_package(
             DynamicExpressions;
-            target_defined_modules=true,
+            target_modules=(DynamicExpressions,),
             ignored_modules=(
                 MyIgnoredModule(DynamicExpressions.NonDifferentiableDeclarationsModule),
             ),
         )
-        # TODO: Hack to get JET to ignore modules
-        # https://github.com/aviatesk/JET.jl/issues/570#issuecomment-2199167755
     end
 end
 if "main" in test_name

test/test_chainrules.jl

@@ -83,9 +83,10 @@ let
 end
 
 # Operator that is NaN for forward pass
-bad_op(x) = x > 0.0 ? log(x) : convert(typeof(x), NaN)
+# Define only for numeric types; `@extend_operators` adds the `Node` method.
+bad_op(x::Real) = x > 0.0 ? log(x) : convert(typeof(x), NaN)
 # And operator that is undefined for backward pass
-undefined_grad_op(x) = x >= 0.0 ? x : zero(x)
+undefined_grad_op(x::Real) = x >= 0.0 ? x : zero(x)
 # And operator that gives a NaN for backward pass
 bad_grad_op(x) = x