Merge pull request #103 from SciML/auto-juliaformatter-pr

Automatic JuliaFormatter.jl run

Author: avik-pal

src/chainrules.jl

@@ -1,5 +1,5 @@
 function CRC.rrule(::Type{<:DeepEquilibriumSolution}, z_star::T, u0::T, residual::T,
-    jacobian_loss::R, nfe::Int) where {T, R <: AbstractFloat}
+        jacobian_loss::R, nfe::Int) where {T, R <: AbstractFloat}
     function ∇deep_equilibrium_solution(dsol)
         return (∂∅, dsol.z_star, dsol.u0, dsol.residual, dsol.jacobian_loss, dsol.nfe)
     end

src/layers/deq.jl

@@ -43,7 +43,7 @@ end
 @truncate_stacktrace DeepEquilibriumNetwork 1 2
 
 function DeepEquilibriumNetwork(model, solver; jacobian_regularization::Bool=false,
-    sensealg=SteadyStateAdjoint(), kwargs...)
+        sensealg=SteadyStateAdjoint(), kwargs...)
     return DeepEquilibriumNetwork{jacobian_regularization}(model, solver, sensealg, kwargs)
 end
 
@@ -112,15 +112,15 @@ end
 @truncate_stacktrace SkipDeepEquilibriumNetwork 1 2 3
 
 function SkipDeepEquilibriumNetwork(model, shortcut, solver; sensealg=SteadyStateAdjoint(),
-    jacobian_regularization::Bool=false, kwargs...)
+        jacobian_regularization::Bool=false, kwargs...)
     return SkipDeepEquilibriumNetwork{jacobian_regularization}(model, shortcut, solver,
         sensealg, kwargs)
 end
 
 _jacobian_regularization(::SkipDeepEquilibriumNetwork{J}) where {J} = J
 
 function _get_initial_condition(deq::SkipDeepEquilibriumNetwork{J, M, Nothing}, x, ps,
-    st) where {J, M}
+        st) where {J, M}
     z, st_ = deq.model((zero(x), x), ps.model, st.model)
     return z, merge(st, (; model=st_))
 end

src/layers/evaluate.jl

@@ -1,5 +1,5 @@
 @generated function _evaluate_unrolled_model(::AbstractDEQs, model, z_star, x, ps, st,
-    ::Val{d}) where {d}
+        ::Val{d}) where {d}
     calls = [:((z_star, st) = model((z_star, x), ps, st)) for _ in 1:d]
     push!(calls, :(return z_star, st))
     return Expr(:block, calls...)

src/layers/jacobian_stabilization.jl

@@ -29,14 +29,14 @@ Estimates the trace of the jacobian matrix wrt `z`.
 Stochastic Estimate of the trace of the Jacobian.
 """
 function estimate_jacobian_trace(::Val{:reverse}, model::Lux.AbstractExplicitLayer,
-    ps, st::NamedTuple, z::AbstractArray, x::AbstractArray, rng::AbstractRNG)
+        ps, st::NamedTuple, z::AbstractArray, x::AbstractArray, rng::AbstractRNG)
     _, back = Zygote.pullback(u -> model((u, x), ps, st)[1], z)
     vjp_z = back(_gaussian_like(rng, x))[1]
     return mean(abs2, vjp_z)
 end
 
 function estimate_jacobian_trace(::Val{:finite_diff}, model::Lux.AbstractExplicitLayer,
-    ps, st::NamedTuple, z::AbstractArray, x::AbstractArray, rng::AbstractRNG)
+        ps, st::NamedTuple, z::AbstractArray, x::AbstractArray, rng::AbstractRNG)
     f = u -> model((u, x), ps, st)[1]
     res = convert(eltype(z), 0)
     epsilon = cbrt(eps(typeof(res)))

src/layers/mdeq.jl

@@ -94,8 +94,8 @@ function Lux.initialstates(rng::AbstractRNG, deq::MultiScaleDeepEquilibriumNetwo
 end
 
 function MultiScaleDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers::Matrix,
-    post_fuse_layer::Union{Nothing, Tuple}, solver, scales::NTuple{N, NTuple{L, Int64}};
-    sensealg=SteadyStateAdjoint(), kwargs...) where {N, L}
+        post_fuse_layer::Union{Nothing, Tuple}, solver, scales::NTuple{N, NTuple{L, Int64}};
+        sensealg=SteadyStateAdjoint(), kwargs...) where {N, L}
     l1 = Parallel(nothing, main_layers...)
     l2 = BranchLayer(Parallel.(+, map(x -> tuple(x...), eachrow(mapping_layers))...)...)
 
@@ -210,7 +210,7 @@ See also: [`DeepEquilibriumNetwork`](@ref), [`SkipDeepEquilibriumNetwork`](@ref)
 end
 
 function MultiScaleSkipDeepEquilibriumNetwork(model::MultiScaleInputLayer{N},
-    args...) where {N}
+        args...) where {N}
     return MultiScaleSkipDeepEquilibriumNetwork{N}(model, args...)
 end
 
@@ -225,9 +225,9 @@ function Lux.initialstates(rng::AbstractRNG, deq::MultiScaleSkipDeepEquilibriumN
 end
 
 function MultiScaleSkipDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers::Matrix,
-    post_fuse_layer::Union{Nothing, Tuple}, shortcut_layers::Union{Nothing, Tuple},
-    solver, scales::NTuple{N, NTuple{L, Int64}};
-    sensealg=SteadyStateAdjoint(), kwargs...) where {N, L}
+        post_fuse_layer::Union{Nothing, Tuple}, shortcut_layers::Union{Nothing, Tuple},
+        solver, scales::NTuple{N, NTuple{L, Int64}};
+        sensealg=SteadyStateAdjoint(), kwargs...) where {N, L}
     l1 = Parallel(nothing, main_layers...)
     l2 = BranchLayer(Parallel.(+, map(x -> tuple(x...), eachrow(mapping_layers))...)...)
     shortcut = shortcut_layers === nothing ? nothing : Parallel(nothing, shortcut_layers...)
@@ -247,7 +247,7 @@ end
 _jacobian_regularization(::MultiScaleSkipDeepEquilibriumNetwork) = false
 
 function _get_initial_condition(deq::MultiScaleSkipDeepEquilibriumNetwork{N, M, Nothing},
-    x, ps, st) where {N, M}
+        x, ps, st) where {N, M}
     u0, st = _get_zeros_initial_condition_mdeq(deq.scales, x, st)
     z, st_ = deq.model((u0, x), ps.model, st.model)
     return z, merge(st, (; model=st_))
@@ -330,8 +330,8 @@ model(x, ps, st)
 See also: [`DeepEquilibriumNetwork`](@ref), [`SkipDeepEquilibriumNetwork`](@ref), [`MultiScaleDeepEquilibriumNetwork`](@ref), [`MultiScaleSkipDeepEquilibriumNetwork`](@ref)
 """
 function MultiScaleNeuralODE(main_layers::Tuple, mapping_layers::Matrix,
-    post_fuse_layer::Union{Nothing, Tuple}, solver, scales::NTuple{N, NTuple{L, Int64}};
-    sensealg=GaussAdjoint(; autojacvec=ZygoteVJP()), kwargs...) where {N, L}
+        post_fuse_layer::Union{Nothing, Tuple}, solver, scales::NTuple{N, NTuple{L, Int64}};
+        sensealg=GaussAdjoint(; autojacvec=ZygoteVJP()), kwargs...) where {N, L}
     l1 = Parallel(nothing, main_layers...)
     l2 = BranchLayer(Parallel.(+, map(x -> tuple(x...), eachrow(mapping_layers))...)...)
 
@@ -361,7 +361,7 @@ end
 
 # Shared Functions
 @generated function _get_zeros_initial_condition_mdeq(::Val{scales}, x::AbstractArray{T, N},
-    st::NamedTuple{fields}) where {scales, T, N, fields}
+        st::NamedTuple{fields}) where {scales, T, N, fields}
     sz = sum(prod.(scales))
     calls = []
     if :initial_condition ∈ fields
@@ -377,6 +377,6 @@ end
 CRC.@non_differentiable _get_zeros_initial_condition_mdeq(::Any...)
 
 @inline function _postprocess_output(deq::Union{MultiScaleDeepEquilibriumNetwork,
-        MultiScaleSkipDeepEquilibriumNetwork, MultiScaleNeuralODE}, z_star)
+            MultiScaleSkipDeepEquilibriumNetwork, MultiScaleNeuralODE}, z_star)
     return split_and_reshape(z_star, deq.split_idxs, deq.scales)
 end

src/solve.jl

@@ -30,8 +30,8 @@ struct ContinuousDEQSolver{A <: DynamicSS} <: AbstractDEQSolver
 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...)
+        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))
@@ -81,7 +81,7 @@ end
 @truncate_stacktrace EquilibriumSolution 1 2
 
 function SciMLBase.__solve(prob::AbstractSteadyStateProblem, alg::AbstractDEQSolver,
-    args...; kwargs...)
+        args...; kwargs...)
     # FIXME: Remove this handle
     sol = solve(prob, alg.alg, args...; kwargs...)
 

src/utils.jl

@@ -27,7 +27,7 @@ DEQs.split_and_reshape(x, split_idxs, shapes)
 ```
 """
 @generated function split_and_reshape(x::AbstractMatrix, ::Val{idxs},
-    ::Val{shapes}) where {idxs, shapes}
+        ::Val{shapes}) where {idxs, shapes}
     dims = [reshape((idxs[i] + 1):idxs[i + 1], shapes[i]...) for i in 1:(length(idxs) - 1)]
     varnames = [gensym("x_view") for _ in dims]
     calls = []

test/adjoint.jl

@@ -15,7 +15,7 @@ function loss_function(model::DEQs.AbstractSkipDeepEquilibriumNetwork, x, ps, st
 end
 
 function loss_function(model::Union{MultiScaleDeepEquilibriumNetwork, MultiScaleNeuralODE},
-    x, ps, st)
+        x, ps, st)
     y, st_ = model(x, ps, st)
     return sum(sum, y) + st_.solution.jacobian_loss
 end