Minor Cleanups

Author: avik-pal

.JuliaFormatter.toml

@@ -2,6 +2,5 @@ style = "sciml"
 whitespace_in_kwargs = false
 always_use_return = true
 format_docstrings = true
-join_lines_based_on_source = false
 separate_kwargs_with_semicolon = true
 format_markdown = true

src/DeepEquilibriumNetworks.jl

@@ -1,22 +1,8 @@
 module DeepEquilibriumNetworks
 
-using CUDA,
-    DiffEqBase,
-    LinearAlgebra,
-    LinearSolve,
-    Lux,
-    MLUtils,
-    OrdinaryDiffEq,
-    Random,
-    SciMLBase,
-    SciMLSensitivity,
-    Setfield,
-    SimpleNonlinearSolve,
-    Static,
-    Statistics,
-    SteadyStateDiffEq,
-    Zygote,
-    ZygoteRules
+using CUDA, DiffEqBase, LinearAlgebra, LinearSolve, Lux, MLUtils, OrdinaryDiffEq, Random,
+    SciMLBase, SciMLSensitivity, Setfield, SimpleNonlinearSolve, Static, Statistics,
+    SteadyStateDiffEq, Zygote, ZygoteRules
 
 using DiffEqBase: AbstractSteadyStateProblem
 using SciMLBase: AbstractNonlinearSolution, AbstractSteadyStateAlgorithm
@@ -48,11 +34,8 @@ export ContinuousDEQSolver, DiscreteDEQSolver
 export EquilibriumSolution, DeepEquilibriumSolution, estimate_jacobian_trace
 
 # Networks
-export DeepEquilibriumNetwork,
-    SkipDeepEquilibriumNetwork,
-    MultiScaleInputLayer,
-    MultiScaleNeuralODE,
-    MultiScaleDeepEquilibriumNetwork,
+export DeepEquilibriumNetwork, SkipDeepEquilibriumNetwork, MultiScaleInputLayer,
+    MultiScaleNeuralODE, MultiScaleDeepEquilibriumNetwork,
     MultiScaleSkipDeepEquilibriumNetwork
 
 end

src/chainrules.jl

@@ -1,15 +1,7 @@
-function CRC.rrule(::Type{<:DeepEquilibriumSolution},
-    z_star::T,
-    u0::T,
-    residual::T,
-    jacobian_loss::R,
-    nfe::Int) where {T, R <: AbstractFloat}
+function CRC.rrule(::Type{<:DeepEquilibriumSolution}, z_star::T, u0::T, residual::T,
+    jacobian_loss::R, nfe::Int) where {T, R <: AbstractFloat}
     function deep_equilibrium_solution_pullback(dsol)
-        return (CRC.NoTangent(),
-            dsol.z_star,
-            dsol.u0,
-            dsol.residual,
-            dsol.jacobian_loss,
+        return (CRC.NoTangent(), dsol.z_star, dsol.u0, dsol.residual, dsol.jacobian_loss,
             dsol.nfe)
     end
     return (DeepEquilibriumSolution(z_star, u0, residual, jacobian_loss, nfe),
@@ -32,9 +24,7 @@ function CRC.rrule(::Type{T}, args...) where {T <: NamedTuple}
     return y, nt_pullback
 end
 
-function CRC.rrule(::typeof(Setfield.set),
-    obj,
-    l::Setfield.PropertyLens{field},
+function CRC.rrule(::typeof(Setfield.set), obj, l::Setfield.PropertyLens{field},
     val) where {field}
     res = Setfield.set(obj, l, val)
     function setfield_pullback(Δres)

src/layers/core.jl

@@ -4,11 +4,8 @@ abstract type AbstractDeepEquilibriumNetwork <:
 function Lux.initialstates(rng::AbstractRNG, deq::AbstractDeepEquilibriumNetwork)
     rng = Lux.replicate(rng)
     randn(rng, 1)
-    return (;
-        model=Lux.initialstates(rng, deq.model),
-        fixed_depth=Val(0),
-        solution=nothing,
-        rng)
+    return (; model=Lux.initialstates(rng, deq.model), fixed_depth=Val(0),
+        solution=nothing, rng)
 end
 
 function Lux.initialparameters(rng::AbstractRNG, deq::AbstractDeepEquilibriumNetwork)
@@ -21,18 +18,13 @@ abstract type AbstractSkipDeepEquilibriumNetwork <:
 function Lux.initialstates(rng::AbstractRNG, deq::AbstractSkipDeepEquilibriumNetwork)
     rng = Lux.replicate(rng)
     randn(rng, 1)
-    return (;
-        model=Lux.initialstates(rng, deq.model),
-        shortcut=Lux.initialstates(rng, deq.shortcut),
-        fixed_depth=Val(0),
-        solution=nothing,
+    return (; model=Lux.initialstates(rng, deq.model),
+        shortcut=Lux.initialstates(rng, deq.shortcut), fixed_depth=Val(0), solution=nothing,
         rng)
 end
 
-const AbstractDEQs = Union{
-    AbstractDeepEquilibriumNetwork,
-    AbstractSkipDeepEquilibriumNetwork,
-}
+const AbstractDEQs = Union{AbstractDeepEquilibriumNetwork,
+    AbstractSkipDeepEquilibriumNetwork}
 
 function (deq::AbstractDEQs)(x::AbstractArray, ps, st::NamedTuple)
     return deq(x, ps, st, _check_unrolled_mode(st))

src/solve.jl

@@ -1,9 +1,9 @@
 abstract type AbstractDEQSolver <: AbstractSteadyStateAlgorithm end
 
 """
-    ContinuousDEQSolver(alg=VCABM3(); mode=SteadyStateTerminationMode.RelSafeBest,
-                        abstol=1.0f-8, reltol=1.0f-6, abstol_termination=abstol,
-                        reltol_termination=reltol, tspan=Inf32, kwargs...)
+    ContinuousDEQSolver(alg=VCAB3(); mode=NLSolveTerminationMode.RelSafeBest,
+        abstol=1.0f-8, reltol=1.0f-6, abstol_termination=abstol, reltol_termination=reltol,
+        tspan=Inf32, kwargs...)
 
 Solver for Continuous DEQ Problem [pal2022mixing](@cite). Effectively a wrapper around
 `DynamicSS` with more sensible defaults for DEQs.
@@ -29,25 +29,18 @@ struct ContinuousDEQSolver{A <: DynamicSS} <: AbstractDEQSolver
     alg::A
 end
 
-function ContinuousDEQSolver(alg=VCAB3();
-    mode=NLSolveTerminationMode.RelSafeBest,
-    abstol=1.0f-8,
-    reltol=1.0f-6,
-    abstol_termination=abstol,
-    reltol_termination=reltol,
-    tspan=Inf32,
-    kwargs...)
-    termination_condition = NLSolveTerminationCondition(mode;
-        abstol=abstol_termination,
-        reltol=reltol_termination,
-        kwargs...)
+function ContinuousDEQSolver(alg=VCAB3(); mode=NLSolveTerminationMode.RelSafeBest,
+    abstol=1.0f-8, reltol=1.0f-6, abstol_termination=abstol, reltol_termination=reltol,
+    tspan=Inf32, kwargs...)
+    termination_condition = NLSolveTerminationCondition(mode; abstol=abstol_termination,
+        reltol=reltol_termination, kwargs...)
     return ContinuousDEQSolver(DynamicSS(alg; abstol, reltol, tspan, termination_condition))
 end
 
 """
     DiscreteDEQSolver(alg = LBroyden(; batched=true,
-                                     termination_condition=NLSolveTerminationCondition(NLSolveTerminationMode.RelSafe;
-                                                                                       abstol=1.0f-8, reltol=1.0f-6))
+        termination_condition=NLSolveTerminationCondition(NLSolveTerminationMode.RelSafe;
+            abstol=1.0f-8, reltol=1.0f-6))
 
 Solver for Discrete DEQ Problem [baideep2019](@cite). Similar to `SSrootfind` but provides
 more flexibility needed for solving DEQ problems.
@@ -60,11 +53,9 @@ See also: [`ContinuousDEQSolver`](@ref)
 """
 Base.@kwdef struct DiscreteDEQSolver{A <: AbstractSimpleNonlinearSolveAlgorithm} <:
                    AbstractDEQSolver
-    alg::A = LBroyden(;
-        batched=true,
+    alg::A = LBroyden(; batched=true,
         termination_condition=NLSolveTerminationCondition(NLSolveTerminationMode.RelSafe;
-            abstol=1.0f-8,
-            reltol=1.0f-6))
+            abstol=1.0f-8, reltol=1.0f-6))
 end
 
 """
@@ -84,10 +75,8 @@ end
 
 @truncate_stacktrace EquilibriumSolution 1 2
 
-function DiffEqBase.__solve(prob::AbstractSteadyStateProblem,
-    alg::AbstractDEQSolver,
-    args...;
-    kwargs...)
+function DiffEqBase.__solve(prob::AbstractSteadyStateProblem, alg::AbstractDEQSolver,
+    args...; kwargs...)
     sol = solve(prob, alg.alg, args...; kwargs...)
 
     # This is not necessarily true and might fail. But makes the code type stable