Add tests for NaN/Inf gradient handling

- Test with custom gradient function that injects NaN periodically
- Test with custom gradient function that injects Inf periodically
- Verify iterations complete and parameters remain finite
- Verify optimizer doesn't crash when encountering bad gradients

Author: ChrisRackauckas

lib/OptimizationOptimisers/test/runtests.jl

@@ -134,3 +134,64 @@ end
 
     @test res.objective < 1e-4
 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
+    x0 = zeros(2)
+    _p = [1.0, 100.0]
+
+    # Counter to track gradient evaluations
+    grad_counter = Ref(0)
+
+    # Custom gradient function that returns NaN on every 5th call
+    function custom_grad!(G, x, p)
+        grad_counter[] += 1
+        if grad_counter[] % 5 == 0
+            # Inject NaN into gradient
+            G .= NaN
+        else
+            # Normal gradient computation
+            G[1] = -2.0 * (p[1] - x[1]) - 4.0 * p[2] * x[1] * (x[2] - x[1]^2)
+            G[2] = 2.0 * p[2] * (x[2] - x[1]^2)
+        end
+        return nothing
+    end
+
+    optprob = OptimizationFunction(rosenbrock; grad = custom_grad!)
+    prob = OptimizationProblem(optprob, x0, _p)
+
+    # Should not throw error and should complete all iterations
+    sol = solve(prob, Optimisers.Adam(0.01), maxiters = 20, progress = false)
+
+    # Verify solution completed all iterations
+    @test sol.stats.iterations == 20
+
+    # 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)
+    function custom_grad_inf!(G, x, p)
+        grad_counter_inf[] += 1
+        if grad_counter_inf[] % 7 == 0
+            # Inject Inf into gradient
+            G .= Inf
+        else
+            # Normal gradient computation
+            G[1] = -2.0 * (p[1] - x[1]) - 4.0 * p[2] * x[1] * (x[2] - x[1]^2)
+            G[2] = 2.0 * p[2] * (x[2] - x[1]^2)
+        end
+        return nothing
+    end
+
+    optprob_inf = OptimizationFunction(rosenbrock; grad = custom_grad_inf!)
+    prob_inf = OptimizationProblem(optprob_inf, x0, _p)
+
+    sol_inf = solve(prob_inf, Optimisers.Adam(0.01), maxiters = 20, progress = false)
+
+    @test sol_inf.stats.iterations == 20
+    @test all(!isnan, sol_inf.u)
+    @test all(isfinite, sol_inf.u)
+end