Fix CI: ChainRules test + JET 0.11 matcher (#153)

Author: MilesCranmerBot

Project.toml

@@ -42,6 +42,8 @@ Reexport = "1"
 SymbolicUtils = "0.19, ^1.0.5, 2, 3"
 Zygote = "0.7"
 julia = "1.10"
+Random = "1"
+TOML = "1"
 
 [extras]
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"

src/ExpressionAlgebra.jl

@@ -35,7 +35,7 @@ function Base.showerror(io::IO, e::MissingOperatorError)
 end
 
 """
-    declare_operator_alias(op::Function, ::Val{arity})::Function
+    declare_operator_alias(op, ::Val{arity})
 
 Define how an internal operator should be matched against user-provided operators in expression trees.
 
@@ -52,7 +52,7 @@ Which would allow a user to write `sqrt(x::Expression)`
 and have it match the operator `safe_sqrt` stored in the binary operators
 of the expression.
 """
-declare_operator_alias(op::F, _) where {F<:Function} = op
+declare_operator_alias(op::F, _) where {F} = op
 
 allow_chaining(@nospecialize(op)) = false
 allow_chaining(::typeof(+)) = true

src/ValueInterface.jl

@@ -24,6 +24,12 @@ Base.@propagate_inbounds function pack_scalar_constants!(
     return idx + 1
 end
 
+# Fallback so callers (and static analysis tools like JET) don't see a MethodError.
+# Types that want to participate in the ValueInterface must implement a more specific method.
+function pack_scalar_constants!(::AbstractVector{<:Number}, ::Int64, value)
+    throw(ArgumentError("pack_scalar_constants! not implemented for $(typeof(value))"))
+end
+
 """
     unpack_scalar_constants(nvals, idx, value)
 
@@ -37,10 +43,14 @@ Returns a tuple of the next index to read from, and the filled-in value.
 """
 Base.@propagate_inbounds function unpack_scalar_constants(
     nvals::AbstractVector{<:Number}, idx::Int64, value::T
-) where {T}
+) where {T<:Number}
     return (idx + 1, convert(T, nvals[idx]))
 end
 
+function unpack_scalar_constants(::AbstractVector{<:Number}, ::Int64, value)
+    throw(ArgumentError("unpack_scalar_constants not implemented for $(typeof(value))"))
+end
+
 """
     count_scalar_constants(value)
 
@@ -60,24 +70,57 @@ end
 function _check_is_valid_array(x)
     return is_valid_array([x]) isa Bool && is_valid_array([x]) == is_valid(x)
 end
-function _check_get_number_type(x)
+function _check_get_number_type(x::X) where {X}
     try
-        get_number_type(typeof(x)) <: Number
+        get_number_type(X) <: Number
     catch e
         @error e
         return false
     end
 end
-function _check_pack_scalar_constants!(x)
-    packed_x = Vector{get_number_type(typeof(x))}(undef, count_scalar_constants(x))
+function _check_pack_scalar_constants!(x::X) where {X}
+    if !applicable(count_scalar_constants, x)
+        return false
+    end
+    n = count_scalar_constants(x)
+
+    packed_x = if X <: Number
+        Vector{X}(undef, n)
+    else
+        # For non-`Number` values, we can't assume a concrete scalar type here.
+        # Use a generic numeric buffer; correctness is checked by roundtripping.
+        Vector{Float64}(undef, n)
+    end
+
+    if !applicable(pack_scalar_constants!, packed_x, 1, x)
+        return false
+    end
+
     new_idx = pack_scalar_constants!(packed_x, 1, x)
-    return new_idx == 1 + count_scalar_constants(x)
+    return new_idx == 1 + n
 end
-function _check_unpack_scalar_constants(x)
-    packed_x = Vector{get_number_type(typeof(x))}(undef, count_scalar_constants(x))
+function _check_unpack_scalar_constants(x::X) where {X}
+    if !applicable(count_scalar_constants, x)
+        return false
+    end
+    n = count_scalar_constants(x)
+
+    packed_x = if X <: Number
+        Vector{X}(undef, n)
+    else
+        Vector{Float64}(undef, n)
+    end
+
+    if !applicable(pack_scalar_constants!, packed_x, 1, x)
+        return false
+    end
+    if !applicable(unpack_scalar_constants, packed_x, 1, x)
+        return false
+    end
+
     pack_scalar_constants!(packed_x, 1, x)
     new_idx, x2 = unpack_scalar_constants(packed_x, 1, x)
-    return new_idx == 1 + count_scalar_constants(x) && x2 == x
+    return new_idx == 1 + n && x2 == x
 end
 function _check_count_scalar_constants(x)
     return count_scalar_constants(x) isa Int &&

src/precompile.jl

@@ -90,7 +90,8 @@ function test_all_combinations(; binary_operators, unary_operators, turbo, types
 end
 
 function test_functions_on_trees(::Type{T}, operators) where {T}
-    local x, c, tree
+    local x, c
+    tree = Node(T; feature=1)
     num_unaops = length(operators.unaops)
     num_binops = length(operators.binops)
     @assert num_unaops > 0 && num_binops > 0

test/Project.toml

@@ -6,8 +6,8 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
+Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
@@ -25,7 +25,8 @@ TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-Aqua = "0.7"
+Aqua = "0.8"
+JET = "0.9, 0.10"
 
 [preferences.DynamicExpressions]
 dispatch_doctor_codegen_level = "min"

test/runtests.jl

@@ -23,26 +23,42 @@ if "jet" in test_name
         )
         using JET
         using DynamicExpressions
-        struct MyIgnoredModule
-            mod::Module
-        end
-        function JET.match_module(
-            mod::MyIgnoredModule, @nospecialize(report::JET.InferenceErrorReport)
-        )
-            s_mod = string(mod.mod)
-            any(report.vst) do vst
-                occursin(s_mod, string(JET.linfomod(vst.linfo)))
+
+        ignored_mod = DynamicExpressions.NonDifferentiableDeclarationsModule
+
+        if isdefined(JET, :match_report) &&
+            isdefined(JET, :ReportMatcher) &&
+            isdefined(JET, :AnyFrameModule)
+            # JET >= 0.11
+            JET.test_package(
+                DynamicExpressions;
+                target_defined_modules=true,
+                ignored_modules=(JET.AnyFrameModule(ignored_mod),),
+            )
+        else
+            # JET <= 0.10: old matcher API
+            struct MyIgnoredModule
+                mod::Module
+            end
+            function JET.match_module(
+                mod::MyIgnoredModule, @nospecialize(report::JET.InferenceErrorReport)
+            )
+                s_mod = string(mod.mod)
+                any(report.vst) do vst
+                    occursin(s_mod, string(JET.linfomod(vst.linfo)))
+                end
             end
+            # On JET 0.10, `target_defined_modules` is not available and also
+            # can cause spurious possible-error reports when analyzing beyond
+            # the package's own modules. Restrict to the DynamicExpressions module.
+            JET.test_package(
+                DynamicExpressions;
+                target_modules=(DynamicExpressions,),
+                ignored_modules=(MyIgnoredModule(ignored_mod),),
+            )
+            # TODO: Hack to get JET to ignore modules
+            # https://github.com/aviatesk/JET.jl/issues/570#issuecomment-2199167755
         end
-        JET.test_package(
-            DynamicExpressions;
-            target_defined_modules=true,
-            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_aqua.jl

@@ -1,4 +1,4 @@
 using DynamicExpressions
 using Aqua
 
-Aqua.test_all(DynamicExpressions; project_toml_formatting=false)
+Aqua.test_all(DynamicExpressions)

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
 
@@ -102,9 +103,11 @@ let
     @extend_operators operators
     x1 = Node(Float64; feature=1)
 
-    nan_forward = bad_op(x1 + 0.5)
-    undefined_grad = undefined_grad_op(x1 + 0.5)
-    nan_grad = bad_grad_op(x1)
+    # `@extend_operators` defines methods via `eval`, which can trigger world-age issues
+    # when tests are executed inside a function (Julia >= 1.12). Use `invokelatest`.
+    nan_forward = Base.invokelatest(bad_op, x1 + 0.5)
+    undefined_grad = Base.invokelatest(undefined_grad_op, x1 + 0.5)
+    nan_grad = Base.invokelatest(bad_grad_op, x1)
 
     function eval_tree(X, tree)
         y, _ = eval_tree_array(tree, X, operators)

test/test_value_interface.jl

@@ -7,7 +7,11 @@ end
 
 @testitem "ValueInterface generic helpers" begin
     using DynamicExpressions.ValueInterfaceModule:
-        is_valid, count_scalar_constants, get_number_type
+        is_valid,
+        count_scalar_constants,
+        get_number_type,
+        pack_scalar_constants!,
+        unpack_scalar_constants
     using Test
 
     # generic `is_valid` (non-number) falls back to `true`
@@ -20,4 +24,19 @@ end
     # simple scalar utilities
     @test count_scalar_constants(42) == 1
     @test get_number_type(Float32) == Float32
+
+    # error paths: fallback methods should throw a helpful error
+    @test_throws ErrorException get_number_type(String)
+    @test_throws ArgumentError pack_scalar_constants!(Float64[], 1, "hello")
+    @test_throws ArgumentError unpack_scalar_constants(Float64[1.0], 1, "hello")
+
+    # internal interface self-checks: verify we handle bad types gracefully
+    @test_logs (:error,) begin
+        @test DynamicExpressions.ValueInterfaceModule._check_get_number_type("hello") ==
+            false
+    end
+    @test DynamicExpressions.ValueInterfaceModule._check_pack_scalar_constants!("hello") ==
+        false
+    @test DynamicExpressions.ValueInterfaceModule._check_unpack_scalar_constants("hello") ==
+        false
 end