Test Nonlinear Solve solvers

Author: avik-pal

src/layers/core.jl

@@ -30,7 +30,7 @@ function (deq::AbstractDEQs)(x::AbstractArray, ps, st::NamedTuple)
 end
 
 # Utilities
-@inline _check_unrolled_mode(::Val{d}) where {d} = Val(d >= 1)
+@inline _check_unrolled_mode(::Val{d}) where {d} = Val(d ≥ 1)
 @inline _check_unrolled_mode(st::NamedTuple) = _check_unrolled_mode(st.fixed_depth)
 
 @inline _get_unrolled_depth(::Val{d}) where {d} = d

src/layers/mdeq.jl

@@ -335,8 +335,7 @@ function MultiScaleNeuralODE(main_layers::Tuple, mapping_layers::Matrix,
         model = MultiScaleInputLayer(Chain(l1, l2), split_idxs, scales)
     else
         model = MultiScaleInputLayer(Chain(l1, l2, Parallel(nothing, post_fuse_layer...)),
-            split_idxs,
-            scales)
+            split_idxs, scales)
     end
 
     return MultiScaleNeuralODE{N}(model, solver, sensealg, scales, split_idxs, kwargs)
@@ -355,8 +354,7 @@ end
 @inline _fix_solution_output(::MultiScaleNeuralODE, x) = x[end]
 
 # Shared Functions
-@generated function _get_zeros_initial_condition_mdeq(::S,
-    x::AbstractArray{T, N},
+@generated function _get_zeros_initial_condition_mdeq(::S, x::AbstractArray{T, N},
     st::NamedTuple{fields}) where {S, T, N, fields}
     scales = known(S)
     sz = sum(prod.(scales))

src/solve.jl

@@ -51,10 +51,16 @@ more flexibility needed for solving DEQ problems.
 
 See also: [`ContinuousDEQSolver`](@ref)
 """
-Base.@kwdef @concrete struct DiscreteDEQSolver <: AbstractDEQSolver
-    alg = LBroyden(; batched=true,
-        termination_condition=NLSolveTerminationCondition(NLSolveTerminationMode.RelSafe;
-            abstol=1.0f-8, reltol=1.0f-6))
+struct DiscreteDEQSolver{A} <: AbstractDEQSolver
+    alg::A
+    function DiscreteDEQSolver(alg=nothing)
+        if alg === nothing
+            alg = LBroyden(; batched=true,
+                termination_condition=NLSolveTerminationCondition(NLSolveTerminationMode.RelSafe;
+                    abstol=1.0f-8, reltol=1.0f-6))
+        end
+        return new{typeof(alg)}(alg)
+    end
 end
 
 """
@@ -78,9 +84,7 @@ function DiffEqBase.__solve(prob::AbstractSteadyStateProblem, alg::AbstractDEQSo
     args...; kwargs...)
     sol = solve(prob, alg.alg, args...; kwargs...)
 
-    # This is not necessarily true and might fail. But makes the code type stable
-    u = sol.u::typeof(prob.u0)
-    du, retcode = sol.resid, sol.retcode
+    u, du, retcode = sol.u, sol.resid, sol.retcode
 
     return EquilibriumSolution{eltype(u), ndims(u)}(u, du, prob, alg, retcode)
 end

test/layers/deq.jl

@@ -122,12 +122,12 @@ function test_skip_deep_equilibrium_network_v2()
     return nothing
 end
 
-Test.@testset "DeepEquilibriumNetwork" begin
+@testset "DeepEquilibriumNetwork" begin
     test_deep_equilibrium_network()
 end
-Test.@testset "SkipDeepEquilibriumNetwork" begin
+@testset "SkipDeepEquilibriumNetwork" begin
     test_skip_deep_equilibrium_network()
 end
-Test.@testset "SkipRegDeepEquilibriumNetwork" begin
+@testset "SkipRegDeepEquilibriumNetwork" begin
     test_skip_deep_equilibrium_network_v2()
 end

test/solve.jl

@@ -4,31 +4,32 @@ using Test
 simple_dudt(u, p, t) = 0.9f0 .* u .- u
 
 function test_continuous_deq_solver()
-    prob = SteadyStateProblem(simple_dudt, [1.0f0], SciMLBase.NullParameters())
+    prob = SteadyStateProblem(simple_dudt, [1.0f0])
 
     sol = solve(prob, ContinuousDEQSolver(); save_everystep=true)
 
     @test sol isa DEQs.EquilibriumSolution
-    @test abs(sol.u[1]) <= 1.0f-4
+    @test abs(sol.u[1]) ≤ 1.0f-4
 
     return nothing
 end
 
-function test_discrete_deq_solver()
-    prob = SteadyStateProblem(simple_dudt, reshape([1.0f0], 1, 1),
-        SciMLBase.NullParameters())
+function test_discrete_deq_solver(; solver=nothing)
+    prob = SteadyStateProblem(simple_dudt, reshape([1.0f0], 1, 1))
 
-    sol = solve(prob, DiscreteDEQSolver())
+    sol = solve(prob, solver === nothing ? DiscreteDEQSolver() : DiscreteDEQSolver(solver))
 
     @test sol isa DEQs.EquilibriumSolution
-    @test abs(sol.u[1]) <= 1.0f-4
+    @test abs(sol.u[1]) ≤ 1.0f-4
 
     return nothing
 end
 
-Test.@testset "Continuous Steady State Solve" begin
+@testset "Continuous Steady State Solve" begin
     test_continuous_deq_solver()
 end
-Test.@testset "Discrete Steady State Solve" begin
-    test_discrete_deq_solver()
+@testset "Discrete Steady State Solve" begin
+    test_discrete_deq_solver(; solver=nothing)  # Default
+    test_discrete_deq_solver(; solver=NewtonRaphson())
+    test_discrete_deq_solver(; solver=LevenbergMarquardt())
 end

test/test_utils.jl

@@ -13,7 +13,7 @@ catch
 end
 
 function get_prng(seed::Int)
-    @static if VERSION >= v"1.7"
+    @static if VERSION ≥ v"1.7"
         rng = Xoshiro()
         Random.seed!(rng, seed)
         return rng
@@ -39,7 +39,7 @@ function is_finite_gradient(gs::NamedTuple)
 end
 
 function run_JET_tests(f, args...; call_broken=false, opt_broken=false, kwargs...)
-    @static if VERSION >= v"1.7"
+    @static if VERSION ≥ v"1.7"
         test_call(f, typeof.(args); broken=call_broken, target_modules=(DEQs,))
         test_opt(f, typeof.(args); broken=opt_broken, target_modules=(DEQs,))
     end