Address PR feedback

- Remove Functors dependency and use simple all(isfinite, G) check
- Make warning conditional on cache.progress flag
- Rewrite tests to use functions that return NaN/Inf in certain regions
  instead of callback-based approach

Author: ChrisRackauckas

lib/OptimizationOptimisers/Project.toml

@@ -4,7 +4,6 @@ authors = ["Vaibhav Dixit <[email protected]> and contributors"]
 version = "0.3.13"
 
 [deps]
-Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -16,7 +15,6 @@ SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 OptimizationBase = {path = "../OptimizationBase"}
 
 [compat]
-Functors = "0.4, 0.5"
 Logging = "1.10"
 Optimisers = "0.2, 0.3, 0.4"
 OptimizationBase = "4"

lib/OptimizationOptimisers/src/OptimizationOptimisers.jl

@@ -3,24 +3,11 @@ module OptimizationOptimisers
 using Reexport, Logging
 @reexport using Optimisers, OptimizationBase
 using SciMLBase
-using Functors
 
 SciMLBase.has_init(opt::AbstractRule) = true
 SciMLBase.requiresgradient(opt::AbstractRule) = true
 SciMLBase.allowsfg(opt::AbstractRule) = true
 
-# Helper function to check if gradients contain NaN or Inf
-function has_nan_or_inf(x)
-    result = Ref(false)
-    Functors.fmap(x) do val
-        if val isa Number && (!isfinite(val))
-            result[] = true
-        end
-        return val
-    end
-    return result[]
-end
-
 function SciMLBase.__init(
         prob::SciMLBase.OptimizationProblem, opt::AbstractRule;
         callback = (args...) -> (false),
@@ -144,9 +131,9 @@ function SciMLBase.__solve(cache::OptimizationCache{O}) where {O <: AbstractRule
             end
         end
         # Skip update if gradient contains NaN or Inf values
-        if !has_nan_or_inf(G)
+        if all(isfinite, G)
             state, θ = Optimisers.update(state, θ, G)
-        else
+        elseif cache.progress
             @warn "Skipping parameter update due to NaN or Inf in gradients at iteration $iterations" maxlog=10
         end
     end

lib/OptimizationOptimisers/test/runtests.jl

@@ -137,54 +137,46 @@ end
 
 @testset "NaN/Inf gradient handling" begin
     # Test that optimizer skips updates when gradients contain NaN or Inf
-    rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+    # Function that returns NaN when parameters are in certain regions
+    function weird_nan_function(x, p)
+        # Return NaN when x[1] is close to certain values to simulate numerical issues
+        if abs(x[1] - 0.3) < 0.05 || abs(x[1] + 0.3) < 0.05
+            return NaN
+        end
+        return (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
+    end
+
     x0 = zeros(2)
     _p = [1.0, 100.0]
 
-    # Test with NaN gradients using Zygote
-    # We'll use a callback to inject NaN into some iterations
-    grad_counter = Ref(0)
-
-    # Create optimization problem with automatic differentiation
-    optprob = OptimizationFunction(rosenbrock, OptimizationBase.AutoZygote())
+    optprob = OptimizationFunction(weird_nan_function, OptimizationBase.AutoZygote())
     prob = OptimizationProblem(optprob, x0, _p)
 
-    # Use a callback that modifies the gradient to inject NaN periodically
-    function nan_callback(state, l)
-        grad_counter[] += 1
-        if grad_counter[] % 5 == 0
-            # Inject NaN into gradient on every 5th iteration
-            state.grad .= NaN
-        end
-        return false
-    end
-
     # Should not throw error and should complete all iterations
-    sol = solve(prob, Optimisers.Adam(0.01), maxiters = 20, progress = false, callback = nan_callback)
+    sol = solve(prob, Optimisers.Adam(0.01), maxiters = 50, progress = false)
 
     # Verify solution completed all iterations
-    @test sol.stats.iterations == 20
+    @test sol.stats.iterations == 50
 
     # Verify parameters are not NaN (would be NaN if updates were applied with NaN gradients)
     @test all(!isnan, sol.u)
     @test all(isfinite, sol.u)
 
-    # Test with Inf gradients
-    grad_counter_inf = Ref(0)
-    prob_inf = OptimizationProblem(optprob, x0, _p)
-
-    function inf_callback(state, l)
-        grad_counter_inf[] += 1
-        if grad_counter_inf[] % 7 == 0
-            # Inject Inf into gradient on every 7th iteration
-            state.grad .= Inf
+    # Function that returns Inf when parameters are in certain regions
+    function weird_inf_function(x, p)
+        # Return Inf when x[1] is close to certain values
+        if abs(x[1] - 0.2) < 0.05 || abs(x[1] + 0.2) < 0.05
+            return Inf
         end
-        return false
+        return (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
     end
 
-    sol_inf = solve(prob_inf, Optimisers.Adam(0.01), maxiters = 20, progress = false, callback = inf_callback)
+    optprob_inf = OptimizationFunction(weird_inf_function, OptimizationBase.AutoZygote())
+    prob_inf = OptimizationProblem(optprob_inf, x0, _p)
+
+    sol_inf = solve(prob_inf, Optimisers.Adam(0.01), maxiters = 50, progress = false)
 
-    @test sol_inf.stats.iterations == 20
+    @test sol_inf.stats.iterations == 50
     @test all(!isnan, sol_inf.u)
     @test all(isfinite, sol_inf.u)
 end