Cleanup codebase

Author: avik-pal

.JuliaFormatter.toml

@@ -3,4 +3,5 @@ whitespace_in_kwargs = false
 always_use_return = true
 format_docstrings = true
 separate_kwargs_with_semicolon = true
-format_markdown = true
+format_markdown = true
+annotate_untyped_fields_with_any = false

.gitignore

@@ -9,4 +9,5 @@ profs
 logs
 benchmarking
 */tensorflow_datasets/
-checkpoints
+checkpoints
+wip

Project.toml

@@ -4,19 +4,20 @@ authors = ["Avik Pal <avikpal@mit.edu>"]
 version = "1.3.0"
 
 [deps]
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
+NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
-SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
@@ -25,17 +26,18 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
-CUDA = "3, 4, 5"
 ChainRulesCore = "1"
+ConcreteStructs = "0.2"
 DiffEqBase = "6.119"
 LinearSolve = "1, 2"
 Lux = "0.4, 0.5"
 MLUtils = "0.2, 0.3, 0.4"
+NonlinearSolve = "2"
 OrdinaryDiffEq = "6"
-SciMLBase = "1.19, 2"
+Reexport = "1"
+SciMLBase = "2"
 SciMLSensitivity = "7"
 Setfield = "1"
-SimpleNonlinearSolve = "0.1.14"
 Static = "0.6, 0.7, 0.8"
 SteadyStateDiffEq = "1.16"
 TruncatedStacktraces = "1.1"

src/DeepEquilibriumNetworks.jl

@@ -1,18 +1,27 @@
 module DeepEquilibriumNetworks
 
-using CUDA, DiffEqBase, LinearAlgebra, LinearSolve, Lux, MLUtils, OrdinaryDiffEq, Random,
-    SciMLBase, SciMLSensitivity, Setfield, SimpleNonlinearSolve, Static, Statistics,
-    SteadyStateDiffEq, Zygote, ZygoteRules
+import Reexport: @reexport
 
-using DiffEqBase: AbstractSteadyStateProblem
-using SciMLBase: AbstractNonlinearSolution, AbstractSteadyStateAlgorithm
-using SimpleNonlinearSolve: AbstractSimpleNonlinearSolveAlgorithm
-using TruncatedStacktraces: @truncate_stacktrace
+@reexport using Lux, NonlinearSolve, OrdinaryDiffEq, SciMLSensitivity
+
+using DiffEqBase, LinearAlgebra, LinearSolve, MLUtils, Random, SciMLBase, SciMLSensitivity,
+    Setfield, Static, Statistics, SteadyStateDiffEq, Zygote, ZygoteRules
+
+import DiffEqBase: AbstractSteadyStateProblem
+import SciMLBase: AbstractNonlinearSolution, AbstractSteadyStateAlgorithm
+import NonlinearSolve: AbstractNonlinearSolveAlgorithm
+import TruncatedStacktraces: @truncate_stacktrace
 
 import ChainRulesCore as CRC
+import ConcreteStructs: @concrete
 
 const DEQs = DeepEquilibriumNetworks
 
+## FIXME: Uses of nothing was removed in Lux 0.5 with a deprecation. It was not updated
+##        here
+Lux.parameterlength(::Nothing) = 0
+Lux.statelength(::Nothing) = 0
+
 include("solve.jl")
 include("utils.jl")
 

src/layers/core.jl

@@ -18,9 +18,8 @@ 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,
-        rng)
+    return (; model=Lux.initialstates(rng, deq.model), rng,
+        shortcut=Lux.initialstates(rng, deq.shortcut), fixed_depth=Val(0), solution=nothing)
 end
 
 const AbstractDEQs = Union{AbstractDeepEquilibriumNetwork,
@@ -54,10 +53,10 @@ Stores the solution of a DeepEquilibriumNetwork and its variants.
     can be computed).
   - `nfe`: Number of Function Evaluations
 """
-struct DeepEquilibriumSolution{T, R <: AbstractFloat, TRes}
+@concrete struct DeepEquilibriumSolution{T, R <: AbstractFloat}
     z_star::T
     u0::T
-    residual::TRes
+    residual
     jacobian_loss::R
-    nfe::Int
+    nfe
 end

src/layers/deq.jl

@@ -22,8 +22,7 @@ using DeepEquilibriumNetworks, Lux, Random, OrdinaryDiffEq
 model = DeepEquilibriumNetwork(Parallel(+,
         Dense(2, 2; use_bias=false),
         Dense(2, 2; use_bias=false)),
-    ContinuousDEQSolver(VCABM3(); abstol=0.01f0, reltol=0.01f0);
-    save_everystep=true)
+    ContinuousDEQSolver(VCABM3(); abstol=0.01f0, reltol=0.01f0); save_everystep=true)
 
 rng = Random.default_rng()
 ps, st = Lux.setup(rng, model)
@@ -34,30 +33,18 @@ model(rand(rng, Float32, 2, 1), ps, st)
 See also: [`SkipDeepEquilibriumNetwork`](@ref), [`MultiScaleDeepEquilibriumNetwork`](@ref),
 [`MultiScaleSkipDeepEquilibriumNetwork`](@ref).
 """
-struct DeepEquilibriumNetwork{J, M, A, S, K} <: AbstractDeepEquilibriumNetwork
-    model::M
-    solver::A
-    sensealg::S
-    kwargs::K
+@concrete struct DeepEquilibriumNetwork{J} <: AbstractDeepEquilibriumNetwork
+    model
+    solver
+    sensealg
+    kwargs
 end
 
 @truncate_stacktrace DeepEquilibriumNetwork 1 2
 
-function DeepEquilibriumNetwork(model,
-    solver;
-    jacobian_regularization::Bool=false,
-    sensealg=SteadyStateAdjoint(),
-    kwargs...)
-    return DeepEquilibriumNetwork{
-        jacobian_regularization,
-        typeof(model),
-        typeof(solver),
-        typeof(sensealg),
-        typeof(kwargs),
-    }(model,
-        solver,
-        sensealg,
-        kwargs)
+function DeepEquilibriumNetwork(model, solver; jacobian_regularization::Bool=false,
+    sensealg=SteadyStateAdjoint(), kwargs...)
+    return DeepEquilibriumNetwork{jacobian_regularization}(model, solver, sensealg, kwargs)
 end
 
 _jacobian_regularization(::DeepEquilibriumNetwork{J}) where {J} = J
@@ -91,8 +78,7 @@ model = SkipDeepEquilibriumNetwork(Parallel(+,
         Dense(2, 2; use_bias=false),
         Dense(2, 2; use_bias=false)),
     Dense(2, 2),
-    ContinuousDEQSolver(VCABM3(); abstol=0.01f0, reltol=0.01f0);
-    save_everystep=true)
+    ContinuousDEQSolver(VCABM3(); abstol=0.01f0, reltol=0.01f0); save_everystep=true)
 
 rng = Random.default_rng()
 ps, st = Lux.setup(rng, model)
@@ -104,8 +90,7 @@ model = SkipDeepEquilibriumNetwork(Parallel(+,
         Dense(2, 2; use_bias=false),
         Dense(2, 2; use_bias=false)),
     nothing,
-    ContinuousDEQSolver(VCABM3(); abstol=0.01f0, reltol=0.01f0);
-    save_everystep=true)
+    ContinuousDEQSolver(VCABM3(); abstol=0.01f0, reltol=0.01f0); save_everystep=true)
 
 rng = Random.default_rng()
 ps, st = Lux.setup(rng, model)
@@ -116,41 +101,25 @@ model(rand(rng, Float32, 2, 1), ps, st)
 See also: [`DeepEquilibriumNetwork`](@ref), [`MultiScaleDeepEquilibriumNetwork`](@ref),
 [`MultiScaleSkipDeepEquilibriumNetwork`](@ref)
 """
-struct SkipDeepEquilibriumNetwork{J, M, Sh, A, S, K} <: AbstractSkipDeepEquilibriumNetwork
-    model::M
-    shortcut::Sh
-    solver::A
-    sensealg::S
-    kwargs::K
+@concrete struct SkipDeepEquilibriumNetwork{J} <: AbstractSkipDeepEquilibriumNetwork
+    model
+    shortcut
+    solver
+    sensealg
+    kwargs
 end
 
 @truncate_stacktrace SkipDeepEquilibriumNetwork 1 2 3
 
-function SkipDeepEquilibriumNetwork(model,
-    shortcut,
-    solver;
-    sensealg=SteadyStateAdjoint(),
-    jacobian_regularization::Bool=false,
-    kwargs...)
-    return SkipDeepEquilibriumNetwork{
-        jacobian_regularization,
-        typeof(model),
-        typeof(shortcut),
-        typeof(solver),
-        typeof(sensealg),
-        typeof(kwargs),
-    }(model,
-        shortcut,
-        solver,
-        sensealg,
-        kwargs)
+function SkipDeepEquilibriumNetwork(model, shortcut, solver; sensealg=SteadyStateAdjoint(),
+    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,
+function _get_initial_condition(deq::SkipDeepEquilibriumNetwork{J, M, Nothing}, x, ps,
     st) where {J, M}
     z, st_ = deq.model((zero(x), x), ps.model, st.model)
     @set! st.model = st_

src/layers/evaluate.jl

@@ -1,9 +1,4 @@
-@generated function _evaluate_unrolled_model(::AbstractDEQs,
-    model,
-    z_star,
-    x,
-    ps,
-    st,
+@generated function _evaluate_unrolled_model(::AbstractDEQs, model, z_star, x, ps, st,
     ::Val{d}) where {d}
     calls = [:((z_star, st) = model((z_star, x), ps, st)) for _ in 1:d]
     push!(calls, :(return z_star, st))
@@ -29,12 +24,7 @@ function (deq::AbstractDEQs)(x::AbstractArray{T}, ps, st::NamedTuple, ::Val{true
     z, st = _get_initial_condition(deq, x, ps, st)
     depth = _get_unrolled_depth(st)
 
-    z_star, st_ = _evaluate_unrolled_model(deq,
-        deq.model,
-        z,
-        x,
-        ps.model,
-        st.model,
+    z_star, st_ = _evaluate_unrolled_model(deq, deq.model, z, x, ps.model, st.model,
         st.fixed_depth)
 
     @set! st.model = st_
@@ -61,13 +51,8 @@ function (deq::AbstractDEQs)(x::AbstractArray, ps, st::NamedTuple, ::Val{false})
 
     if _jacobian_regularization(deq)
         rng = Lux.replicate(st.rng)
-        jac_loss = estimate_jacobian_trace(Val(:finite_diff),
-            deq.model,
-            ps.model,
-            st.model,
-            z_star,
-            x,
-            rng)
+        jac_loss = estimate_jacobian_trace(Val(:finite_diff), deq.model, ps.model, st.model,
+            z_star, x, rng)
     else
         rng = st.rng
         jac_loss = T(0)

src/layers/jacobian_stabilization.jl

@@ -1,5 +1,10 @@
-_gaussian_like(rng::AbstractRNG, x::AbstractArray) = randn(rng, eltype(x), size(x))
-_gaussian_like(rng::AbstractRNG, x::CuArray) = CUDA.randn(rng, eltype(x), size(x))
+function _gaussian_like(rng::AbstractRNG, x)
+    y = similar(x)
+    randn!(rng, y)
+    return y
+end
+
+CRC.@non_differentiable _gaussian_like(::Any...)
 
 """
     estimate_jacobian_trace(::Val{mode}, model::Lux.AbstractExplicitLayer, ps,
@@ -23,25 +28,15 @@ 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)
+function estimate_jacobian_trace(::Val{:reverse}, model::Lux.AbstractExplicitLayer,
+    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)
+function estimate_jacobian_trace(::Val{:finite_diff}, model::Lux.AbstractExplicitLayer,
+    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)))