Clean up the codebase and use the latest versions of NonlinearSolve and SteadyStateDiffEq

Author: avik-pal

.JuliaFormatter.toml

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

.github/workflows/CI.yml

@@ -21,10 +21,6 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        group:
-          - CORE
-          - LAYERS
-          - ADJOINT
         version:
           - '1'
     steps:
@@ -44,9 +40,9 @@ jobs:
             ${{ runner.os }}-
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
-        env:
-          GROUP: ${{ matrix.group }}
       - uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: src,ext
       - uses: codecov/codecov-action@v3
         with:
           files: lcov.info

Project.toml

@@ -1,43 +1,43 @@
 name = "DeepEquilibriumNetworks"
 uuid = "6748aba7-0e9b-415e-a410-ae3cc0ecb334"
 authors = ["Avik Pal <avikpal@mit.edu>"]
-version = "1.4.0"
+version = "2.0.0"
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
-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"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
+
+[weakdeps]
+SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[extensions]
+DeepEquilibriumNetworksSciMLSensitivityExt = "SciMLSensitivity"
+DeepEquilibriumNetworksZygoteExt = "Zygote"
+
 [compat]
+ADTypes = "0.2.5"
 ChainRulesCore = "1"
 ConcreteStructs = "0.2"
+ConstructionBase = "1"
 DiffEqBase = "6.119"
 LinearAlgebra = "1"
-LinearSolve = "1, 2"
 Lux = "0.5.7"
-NonlinearSolve = "2"
-OrdinaryDiffEq = "6"
 Random = "1"
-Reexport = "1"
 SciMLBase = "2"
 SciMLSensitivity = "7.43"
-Setfield = "1"
 Statistics = "1"
-SteadyStateDiffEq = "1.16"
+SteadyStateDiffEq = "2"
 TruncatedStacktraces = "1.1"
-Zygote = "0.6.34"
+Zygote = "0.6.67"
 julia = "1.9"

README.md

@@ -57,14 +57,11 @@ If you are using this project for research or other academic purposes consider c
 paper:
 
 ```bibtex
-@misc{pal2022mixing,
-  title={Mixing Implicit and Explicit Deep Learning with Skip DEQs and Infinite Time Neural
-         ODEs (Continuous DEQs)}, 
-  author={Avik Pal and Alan Edelman and Christopher Rackauckas},
-  year={2022},
-  eprint={2201.12240},
-  archivePrefix={arXiv},
-  primaryClass={cs.LG}
+@article{pal2022continuous,
+  title={Continuous Deep Equilibrium Models: Training Neural ODEs Faster by Integrating Them to Infinity},
+  author={Pal, Avik and Edelman, Alan and Rackauckas, Christopher},
+  booktitle={2023 IEEE High Performance Extreme Computing Conference (HPEC)}, 
+  year={2023}
 }
 ```
 

docs/Project.toml

@@ -4,6 +4,6 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 
 [compat]
-DeepEquilibriumNetworks = "1"
+DeepEquilibriumNetworks = "2"
 Documenter = "1"
 DocumenterCitations = "0.2, 1"

docs/make.jl

@@ -7,17 +7,10 @@ bib = CitationBibliography(joinpath(@__DIR__, "ref.bib"); style=:authoryear)
 
 include("pages.jl")
 
-makedocs(;
-    sitename="Fast Deep Equilibrium Networks",
-    authors="Avik Pal et al.",
-    modules=[DeepEquilibriumNetworks],
-    clean=true,
-    doctest=false,
-    linkcheck=true,
-    format=Documenter.HTML(;
-        assets=["assets/favicon.ico"],
+makedocs(; sitename="Deep Equilibrium Networks", authors="Avik Pal et al.",
+    modules=[DeepEquilibriumNetworks], clean=true, doctest=true,
+    linkcheck=true, format=Documenter.HTML(; assets=["assets/favicon.ico"],
         canonical="https://docs.sciml.ai/DeepEquilibriumNetworks/stable/"),
-    plugins=[bib],
-    pages)
+    plugins=[bib], pages)
 
 deploydocs(; repo="github.com/SciML/DeepEquilibriumNetworks.jl.git", push_preview=true)

docs/pages.jl

@@ -1,9 +1,6 @@
 pages = [
     "Home" => "index.md",
-    "Manual" => [
-        "Dynamical Systems" => "manual/solvers.md",
-        "DEQ Layers" => "manual/deqs.md",
-        "Miscellaneous" => "manual/misc.md",
-    ],
+    "Tutorial" => "tutorial.md",
+    "API References" => "api.md",
     "References" => "references.md",
 ]

docs/src/api.md

@@ -0,0 +1,31 @@
+# Deep Equilibrium Models
+
+[baideep2019](@cite) introduced Discrete Deep Equilibrium Models which drives a Discrete
+Dynamical System to its steady-state. [pal2022mixing](@cite) extends this framework to
+Continuous Dynamical Systems which converge to the steady-stable in a more stable fashion.
+For a detailed discussion refer to [pal2022mixing](@cite).
+
+To construct a continuous DEQ, any ODE solver compatible with `DifferentialEquations.jl` API
+can be passed as the solver. To construct a discrete DEQ, any root finding algorithm
+compatible with `NonlinearSolve.jl` API can be passed as the solver.
+
+## Standard Models
+
+```@docs
+DeepEquilibriumNetwork
+SkipDeepEquilibriumNetwork
+```
+
+## MultiScale Models
+
+```@docs
+MultiScaleDeepEquilibriumNetwork
+MultiScaleSkipDeepEquilibriumNetwork
+MultiScaleNeuralODE
+```
+
+## Solution
+
+```@docs
+DeepEquilibriumSolution
+```