Rework package structure and tests

.github/workflows/CI.yml

@@ -15,8 +15,14 @@ jobs:
   test:
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
+        group:
+          - CORE
+          - LAYERS
+          - ADJOINT
         version:
+          - 1.6
           - 1.7
     steps:
       - uses: actions/checkout@v2
@@ -35,6 +41,8 @@ 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
       - uses: codecov/codecov-action@v2
         with:

.github/workflows/Documentation.yml

@@ -16,10 +16,16 @@ jobs:
         with:
           version: "1"
       - name: Install dependencies
-        # FIXME: Remove once Lux.jl is registered
-        run: julia --project=docs -e 'using Pkg; Pkg.add(url="https://github.com/avik-pal/Lux.jl");  Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+        run: julia --project=docs -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token
           DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} # For authentication with SSH deploy key
-        run: julia --project=docs/ docs/make.jl
+          GKSwstype: "100" # https://discourse.julialang.org/t/generation-of-documentation-fails-qt-qpa-xcb-could-not-connect-to-display/60988
+          JULIA_DEBUG: "Documenter"
+          DATADEPS_ALWAYS_ACCEPT: true
+        run: julia --code-coverage=user --project=docs/ --color=yes docs/make.jl
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v2
+        with:
+          files: lcov.info

.github/workflows/FormatCheck.yml

@@ -0,0 +1,43 @@
+name: format-check
+
+on:
+  push:
+    branches:
+      - 'main'
+      - 'release-'
+    tags: '*'
+  pull_request:
+
+jobs:
+  build:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        julia-version: [1.7]
+        julia-arch: [x86]
+        os: [ubuntu-latest]
+    steps:
+      - uses: julia-actions/setup-julia@latest
+        with:
+          version: ${{ matrix.julia-version }}
+
+      - uses: actions/checkout@v1
+      - name: Install JuliaFormatter and format
+        # This will use the latest version by default but you can set the version like so:
+        #
+        # julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="0.13.0"))'
+        run: |
+          julia  -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))'
+          julia  -e 'using JuliaFormatter; format(".", verbose=true)'
+      - name: Format check
+        run: |
+          julia -e '
+          out = Cmd(`git diff --name-only`) |> read |> String
+          if out == ""
+              exit(0)
+          else
+              @error "Some files have not been formatted !!!"
+              write(stdout, out)
+              exit(1)
+          end'
+      

Project.toml

@@ -1,12 +1,11 @@
 name = "DeepEquilibriumNetworks"
 uuid = "6748aba7-0e9b-415e-a410-ae3cc0ecb334"
 authors = ["Avik Pal <[email protected]>"]
-version = "0.1.2"
+version = "0.2.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
@@ -18,7 +17,6 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
-Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
@@ -28,7 +26,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [compat]
 CUDA = "3"
 ChainRulesCore = "1"
-ComponentArrays = "0.11, 0.12"
 DiffEqBase = "6"
 DiffEqCallbacks = "2.20.1"
 Functors = "0.2, 0.3"
@@ -38,9 +35,8 @@ MLUtils = "0.2"
 OrdinaryDiffEq = "6"
 SciMLBase = "1.19"
 SciMLSensitivity = "7"
-Setfield = "1"
 Static = "0.6, 0.7"
 SteadyStateDiffEq = "1.6"
 UnPack = "1"
 Zygote = "0.6.34"
-julia = "1.7"
+julia = "1.6"

README.md

@@ -2,7 +2,66 @@
 
 [![Latest Docs](https://img.shields.io/badge/docs-latest-blue.svg)](https://deepequilibriumnetworks.sciml.ai/dev/)
 [![Stable Docs](https://img.shields.io/badge/docs-stable-blue.svg)](https://deepequilibriumnetworks.sciml.ai/stable/)
+[![CI](https://github.com/SciML/DeepEquilibriumNetworks.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/SciML/DeepEquilibriumNetworks.jl/actions/workflows/CI.yml)
 [![codecov](https://codecov.io/gh/SciML/DeepEquilibriumNetworks.jl/branch/main/graph/badge.svg?token=plksEh6pUG)](https://codecov.io/gh/SciML/DeepEquilibriumNetworks.jl)
 [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
+[![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle)
+[![Package Downloads](https://shields.io/endpoint?url=https://pkgs.genieframework.com/api/v1/badge/DeepEquilibriumNetworks)](https://pkgs.genieframework.com?packages=DeepEquilibriumNetworks)
 
-Deep Equilibrium Networks using [Lux.jl](https://lux.csail.mit.edu/dev) and [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/)
+DeepEquilibriumNetworks.jl is a framework built on top of
+[DifferentialEquations.jl](https://diffeq.sciml.ai/stable/) and
+[Lux.jl](https://lux.csail.mit.edu/dev/) enabling the efficient training and inference for
+Deep Equilibrium Networks (Infinitely Deep Neural Networks).
+
+## Installation
+
+```julia
+] add DeepEquilibriumNetworks
+```
+
+## Quickstart
+
+```julia
+import DeepEquilibriumNetworks as DEQs
+import Lux
+import Random
+import Zygote
+
+seed = 0
+rng = Random.default_rng()
+Random.seed!(rng, seed)
+
+model = Lux.Chain(Lux.Dense(2, 2),
+                  DEQs.DeepEquilibriumNetwork(Lux.Parallel(+, Lux.Dense(2, 2; bias=false),
+                                                           Lux.Dense(2, 2; bias=false)),
+                                              DEQs.ContinuousDEQSolver(; abstol=0.1f0,
+                                                                       reltol=0.1f0,
+                                                                       abstol_termination=0.1f0,
+                                                                       reltol_termination=0.1f0)))
+
+ps, st = gpu.(Lux.setup(rng, model))
+x = gpu(rand(rng, Float32, 2, 1))
+y = gpu(rand(rng, Float32, 2, 1))
+
+gs = Zygote.gradient(p -> sum(abs2, model(x, p, st)[1][1] .- y), ps)[1]
+```
+
+## Citation
+
+If you are using this project for research or other academic purposes consider citing our
+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}
+}
+```
+
+For specific algorithms, check the respective documentations and cite the corresponding
+papers.

docs/make.jl

@@ -4,16 +4,22 @@ bib = CitationBibliography(joinpath(@__DIR__, "ref.bib"); sorting=:nyt)
 
 makedocs(bib; sitename="Fast Deep Equilibrium Networks", authors="Avik Pal et al.",
          clean=true, doctest=false, modules=[DeepEquilibriumNetworks],
-         format=Documenter.HTML(#  analytics = "",
-                                #  assets = ["assets/favicon.ico"],
-                                ;
+         strict=[
+           :doctest,
+           :linkcheck,
+           :parse_error,
+           :example_block,
+           # Other available options are
+           # :autodocs_block, :cross_references, :docs_block, :eval_block, :example_block,
+           # :footnote, :meta_block, :missing_docs, :setup_block
+         ], checkdocs=:all,
+         format=Documenter.HTML(;
                                 canonical="https://deepequilibriumnetworks.sciml.ai/stable/"),
          pages=[
            "Home" => "index.md",
            "Manual" => [
              "Dynamical Systems" => "manual/solvers.md",
              "Non Linear Solvers" => "manual/nlsolve.md",
-             "General Purpose Layers" => "manual/layers.md",
              "DEQ Layers" => "manual/deqs.md",
              "Miscellaneous" => "manual/misc.md",
            ],

docs/src/index.md

@@ -1,6 +1,17 @@
 # DeepEquilibriumNetworks: (Fast) Deep Equlibrium Networks
 
-DeepEquilibriumNetworks.jl is a framework built on top of [DifferentialEquations.jl](https://diffeq.sciml.ai/stable/) and [Lux.jl](https://lux.csail.mit.edu/dev/) enabling the efficient training and inference for Deep Equilibrium Networks (Infinitely Deep Neural Networks).
+[![Latest Docs](https://img.shields.io/badge/docs-latest-blue.svg)](https://deepequilibriumnetworks.sciml.ai/dev/)
+[![Stable Docs](https://img.shields.io/badge/docs-stable-blue.svg)](https://deepequilibriumnetworks.sciml.ai/stable/)
+[![CI](https://github.com/SciML/DeepEquilibriumNetworks.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/SciML/DeepEquilibriumNetworks.jl/actions/workflows/CI.yml)
+[![codecov](https://codecov.io/gh/SciML/DeepEquilibriumNetworks.jl/branch/main/graph/badge.svg?token=plksEh6pUG)](https://codecov.io/gh/SciML/DeepEquilibriumNetworks.jl)
+[![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
+[![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle)
+[![Package Downloads](https://shields.io/endpoint?url=https://pkgs.genieframework.com/api/v1/badge/DeepEquilibriumNetworks)](https://pkgs.genieframework.com?packages=DeepEquilibriumNetworks)
+
+DeepEquilibriumNetworks.jl is a framework built on top of
+[DifferentialEquations.jl](https://diffeq.sciml.ai/stable/) and
+[Lux.jl](https://lux.csail.mit.edu/dev/) enabling the efficient training and inference for
+Deep Equilibrium Networks (Infinitely Deep Neural Networks).
 
 ## Installation
 
@@ -11,40 +22,46 @@ DeepEquilibriumNetworks.jl is a framework built on top of [DifferentialEquations
 ## Quickstart
 
 ```julia
-using Lux, DeepEquilibriumNetworks, Random
+import DeepEquilibriumNetworks as DEQs
+import Lux
+import Random
+import Zygote
 
 seed = 0
 rng = Random.default_rng()
 Random.seed!(rng, seed)
 
-model = DEQChain(Dense(2, 2),
-                 DeepEquilibriumNetwork(Parallel(+, Dense(2, 2; bias=false),
-                                                 Dense(2, 2; bias=false)),
-                                        ContinuousDEQSolver(; abstol=0.1f0,
-                                                            reltol=0.1f0,
-                                                            abstol_termination=0.1f0,
-                                                            reltol_termination=0.1f0)))
+model = Lux.Chain(Lux.Dense(2, 2),
+                  DEQs.DeepEquilibriumNetwork(Lux.Parallel(+, Lux.Dense(2, 2; bias=false),
+                                                           Lux.Dense(2, 2; bias=false)),
+                                              DEQs.ContinuousDEQSolver(; abstol=0.1f0,
+                                                                       reltol=0.1f0,
+                                                                       abstol_termination=0.1f0,
+                                                                       reltol_termination=0.1f0)))
 
 ps, st = gpu.(Lux.setup(rng, model))
 x = gpu(rand(rng, Float32, 2, 1))
 y = gpu(rand(rng, Float32, 2, 1))
 
-gs = gradient(p -> sum(abs2, model(x, p, st)[1][1] .- y), ps)[1]
+gs = Zygote.gradient(p -> sum(abs2, model(x, p, st)[1][1] .- y), ps)[1]
 ```
 
 ## Citation
 
-If you are using this project for research or other academic purposes consider citing our paper:
+If you are using this project for research or other academic purposes consider citing our
+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}
+  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}
 }
 ```
 
-For specific algorithms, check the respective documentations and cite the corresponding papers.
+For specific algorithms, check the respective documentations and cite the corresponding
+papers.

docs/src/manual/misc.md

@@ -3,5 +3,4 @@
 ```@docs
 DeepEquilibriumAdjoint
 DeepEquilibriumSolution
-NormalInitializer
 ```

src/DeepEquilibriumNetworks.jl

@@ -1,18 +1,10 @@
 module DeepEquilibriumNetworks
 
-using ChainRulesCore, ComponentArrays, CUDA, DiffEqBase, DiffEqCallbacks, SciMLSensitivity,
-      Functors, LinearAlgebra, LinearSolve, Lux, MLUtils, OrdinaryDiffEq, SciMLBase,
-      Setfield, Static, Statistics, SteadyStateDiffEq, UnPack, Zygote
-
-import SciMLSensitivity: AbstractAdjointSensitivityAlgorithm
-import Lux: AbstractExplicitContainerLayer, initialparameters, initialstates,
-            parameterlength, statelength
-import Random: AbstractRNG
-
 include("operator.jl")
 
-include("solvers/continuous.jl")
-include("solvers/discrete.jl")
+include("solvers/solvers.jl")
+include("solvers/discrete/broyden.jl")
+include("solvers/discrete/limited_memory_broyden.jl")
 include("solvers/termination.jl")
 
 include("solve.jl")
@@ -22,18 +14,17 @@ include("layers/core.jl")
 include("layers/jacobian_stabilization.jl")
 include("layers/deq.jl")
 include("layers/mdeq.jl")
-include("layers/chain.jl")
 
 include("adjoint.jl")
 
 # DEQ Solvers
 export ContinuousDEQSolver, DiscreteDEQSolver, BroydenSolver, LimitedMemoryBroydenSolver
 
 # Utils
-export NormalInitializer, DeepEquilibriumAdjoint, compute_deq_jacobian_loss,
-       DeepEquilibriumSolution
+export DeepEquilibriumAdjoint, DeepEquilibriumSolution, estimate_jacobian_trace
 
-export DeepEquilibriumNetwork, SkipDeepEquilibriumNetwork, MultiScaleDeepEquilibriumNetwork,
-       MultiScaleSkipDeepEquilibriumNetwork, DEQChain
+# Networks
+export DeepEquilibriumNetwork, SkipDeepEquilibriumNetwork
+export MultiScaleDeepEquilibriumNetwork, MultiScaleSkipDeepEquilibriumNetwork
 
 end