Merge pull request #146 from SciML/ap/fix_tests

Retrigger Tests

Author: avik-pal

.github/workflows/CI.yml

@@ -23,7 +23,6 @@ jobs:
       matrix:
         version:
           - '1'
-          - '~1.10.0-0'
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1

Project.toml

@@ -1,16 +1,18 @@
 name = "DeepEquilibriumNetworks"
 uuid = "6748aba7-0e9b-415e-a410-ae3cc0ecb334"
 authors = ["Avik Pal <[email protected]>"]
-version = "2.0.2"
+version = "2.0.3"
 
 [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"
+FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -32,9 +34,11 @@ ChainRulesCore = "1"
 ConcreteStructs = "0.2"
 ConstructionBase = "1"
 DiffEqBase = "6.119"
+FastClosures = "0.3"
 LinearAlgebra = "1"
 LinearSolve = "2.21.2"
 Lux = "0.5.11"
+PrecompileTools = "1"
 Random = "1"
 SciMLBase = "2"
 SciMLSensitivity = "7.43"

docs/Project.toml

@@ -3,6 +3,7 @@ DeepEquilibriumNetworks = "6748aba7-0e9b-415e-a410-ae3cc0ecb334"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLDataUtils = "cc2ba9b6-d476-5e6d-8eaf-a92d5412d41d"
@@ -20,6 +21,7 @@ DeepEquilibriumNetworks = "2"
 Documenter = "1"
 DocumenterCitations = "1"
 LinearSolve = "2"
+LoggingExtras = "1"
 Lux = "0.5"
 LuxCUDA = "0.3"
 MLDataUtils = "0.5"

docs/ref.bib

@@ -19,7 +19,6 @@ @article{baideep2019
 	year = {2019},
 	note = {arXiv: 1909.01377},
 	keywords = {Statistics - Machine Learning, Computer Science - Machine Learning},
-	annote = {Comment: NeurIPS 2019 Spotlight Oral},
 	file = {Bai et al. - 2019 - Deep Equilibrium Models.pdf:files/245/Bai et al. - 2019 - Deep Equilibrium Models.pdf:application/pdf},
 }
 
@@ -35,7 +34,6 @@ @article{baimultiscale2020
 	year = {2020},
 	note = {arXiv: 2006.08656},
 	keywords = {Statistics - Machine Learning, Computer Science - Machine Learning, Computer Science - Computer Vision and Pattern Recognition},
-	annote = {Comment: NeurIPS 2020 Oral},
 	file = {Bai et al. - 2020 - Multiscale Deep Equilibrium Models.pdf:files/248/Bai et al. - 2020 - Multiscale Deep Equilibrium Models.pdf:application/pdf},
 }
 

docs/src/tutorials/basic_mnist_deq.md

@@ -4,7 +4,7 @@ We will train a simple Deep Equilibrium Model on MNIST. First we load a few pack
 
 ```@example basic_mnist_deq
 using DeepEquilibriumNetworks, SciMLSensitivity, Lux, NonlinearSolve, OrdinaryDiffEq,
-    Statistics, Random, Optimisers, LuxCUDA, Zygote, LinearSolve
+    Statistics, Random, Optimisers, LuxCUDA, Zygote, LinearSolve, LoggingExtras
 using MLDatasets: MNIST
 using MLDataUtils: LabelEnc, convertlabel, stratifiedobs, batchview
 
@@ -20,6 +20,18 @@ const cdev = cpu_device()
 const gdev = gpu_device()
 ```
 
+SciMLBase introduced a warning instead of depwarn which pollutes the output. We can suppress
+it with the following logger
+
+```@example basic_mnist_deq
+function remove_syms_warning(log_args)
+    return log_args.message !=
+           "The use of keyword arguments `syms`, `paramsyms` and `indepsym` for `SciMLFunction`s is deprecated. Pass `sys = SymbolCache(syms, paramsyms, indepsym)` instead."
+end
+
+filtered_logger = ActiveFilteredLogger(remove_syms_warning, global_logger())
+```
+
 We can now construct our dataloader.
 
 ```@example basic_mnist_deq
@@ -175,15 +187,19 @@ and end up using solvers like `Broyden`, but we can simply slap in any of the fa
 from NonlinearSolve.jl. Here we will use Newton-Krylov Method:
 
 ```@example basic_mnist_deq
-train_model(NewtonRaphson(; linsolve=KrylovJL_GMRES()), :regdeq)
+with_logger(filtered_logger) do
+    train_model(NewtonRaphson(; linsolve=KrylovJL_GMRES()), :regdeq)
+end
 nothing # hide
 ```
 
 We can also train a continuous DEQ by passing in an ODE solver. Here we will use `VCAB3()`
 which tend to be quite fast for continuous Neural Network problems.
 
 ```@example basic_mnist_deq
-train_model(VCAB3(), :deq)
+with_logger(filtered_logger) do
+    train_model(VCAB3(), :deq)
+end
 nothing # hide
 ```
 

docs/src/tutorials/reduced_dim_deq.md

@@ -6,7 +6,7 @@ same MNIST example as before, but this time we will use a reduced state size.
 
 ```@example reduced_dim_mnist
 using DeepEquilibriumNetworks, SciMLSensitivity, Lux, NonlinearSolve, OrdinaryDiffEq,
-    Statistics, Random, Optimisers, LuxCUDA, Zygote, LinearSolve
+    Statistics, Random, Optimisers, LuxCUDA, Zygote, LinearSolve, LoggingExtras
 using MLDatasets: MNIST
 using MLDataUtils: LabelEnc, convertlabel, stratifiedobs, batchview
 
@@ -16,6 +16,13 @@ ENV["DATADEPS_ALWAYS_ACCEPT"] = true
 const cdev = cpu_device()
 const gdev = gpu_device()
 
+function remove_syms_warning(log_args)
+    return log_args.message !=
+           "The use of keyword arguments `syms`, `paramsyms` and `indepsym` for `SciMLFunction`s is deprecated. Pass `sys = SymbolCache(syms, paramsyms, indepsym)` instead."
+end
+
+filtered_logger = ActiveFilteredLogger(remove_syms_warning, global_logger())
+
 function onehot(labels_raw)
     return convertlabel(LabelEnc.OneOfK, labels_raw, LabelEnc.NativeLabels(collect(0:9)))
 end
@@ -168,11 +175,15 @@ Now we can train our model. We can't use `:regdeq` here currently, but we will s
 in the future.
 
 ```@example reduced_dim_mnist
-train_model(NewtonRaphson(; linsolve=KrylovJL_GMRES()), :skipdeq)
+with_logger(filtered_logger) do
+    train_model(NewtonRaphson(; linsolve=KrylovJL_GMRES()), :skipdeq)
+end
 nothing # hide
 ```
 
 ```@example reduced_dim_mnist
-train_model(NewtonRaphson(; linsolve=KrylovJL_GMRES()), :deq)
+with_logger(filtered_logger) do
+    train_model(NewtonRaphson(; linsolve=KrylovJL_GMRES()), :deq)
+end
 nothing # hide
 ```

src/DeepEquilibriumNetworks.jl

@@ -1,16 +1,19 @@
 module DeepEquilibriumNetworks
 
-using ADTypes,
-    DiffEqBase, LinearAlgebra, Lux, Random, SciMLBase, Statistics, SteadyStateDiffEq
+import PrecompileTools: @recompile_invalidations
 
-import ChainRulesCore as CRC
-import ConcreteStructs: @concrete
-import ConstructionBase: constructorof
-import Lux: AbstractExplicitLayer, AbstractExplicitContainerLayer
-import TruncatedStacktraces: @truncate_stacktrace
+@recompile_invalidations begin
+    using ADTypes, DiffEqBase, FastClosures, LinearAlgebra, Lux, Random, SciMLBase,
+        Statistics, SteadyStateDiffEq
 
-import SciMLBase: AbstractNonlinearAlgorithm,
-    AbstractODEAlgorithm, _unwrap_val, NonlinearSolution
+    import ChainRulesCore as CRC
+    import ConcreteStructs: @concrete
+    import ConstructionBase: constructorof
+    import Lux: AbstractExplicitLayer, AbstractExplicitContainerLayer
+    import SciMLBase: AbstractNonlinearAlgorithm,
+        AbstractODEAlgorithm, _unwrap_val, NonlinearSolution
+    import TruncatedStacktraces: @truncate_stacktrace
+end
 
 # Useful Constants
 const DEQs = DeepEquilibriumNetworks

src/layers.jl

@@ -97,7 +97,12 @@ function (deq::DEQ{pType})(x, ps, st::NamedTuple, ::Val{false}) where {pType}
 
     model = Lux.Experimental.StatefulLuxLayer(deq.model, nothing, st.model)
 
-    dudt(u, p, t) = model((u, p.x), p.ps) .- u
+    dudt = @closure (u, p, t) -> begin
+        # The type-assert is needed because of an upstream Lux issue with type stability of
+        # conv with Dual numbers
+        y = model((u, p.x), p.ps)::typeof(u)
+        return y .- u
+    end
 
     prob = __construct_prob(pType, ODEFunction{false}(dudt), z, (; ps=ps.model, x))
     alg = __normalize_alg(deq)
@@ -144,7 +149,7 @@ Deep Equilibrium Network as proposed in [baideep2019](@cite) and [pal2022mixing]
 
 ## Example
 
-```jldoctest
+```julia
 julia> using DeepEquilibriumNetworks, Lux, Random, OrdinaryDiffEq
 
 julia> model = DeepEquilibriumNetwork(Parallel(+, Dense(2, 2; use_bias=false),
@@ -225,7 +230,7 @@ For keyword arguments, see [`DeepEquilibriumNetwork`](@ref).
 
 ## Example
 
-```jldoctest
+```julia
 julia> using DeepEquilibriumNetworks, Lux, Random, NonlinearSolve
 
 julia> main_layers = (Parallel(+, Dense(4 => 4, tanh; use_bias=false),

test/Project.toml

@@ -19,6 +19,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TestSetExtensions = "98d24dd4-01ad-11ea-1b02-c9a08f80db04"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]

test/layers.jl

@@ -22,7 +22,9 @@ end
     x_sizes = [(2, 14), (3, 3, 1, 3)]
 
     model_type = (:deq, :skipdeq, :skipregdeq)
-    solvers = (VCAB3(), Tsit5(), NewtonRaphson(), SimpleLimitedMemoryBroyden())
+    solvers = (VCAB3(), Tsit5(),
+        NewtonRaphson(; autodiff=AutoForwardDiff(; chunksize=12)),
+        SimpleLimitedMemoryBroyden())
     jacobian_regularizations = Any[nothing, AutoZygote()]
     !ongpu && push!(jacobian_regularizations, AutoFiniteDiff())
 
@@ -31,8 +33,6 @@ end
 
         @testset "x_size: $(x_size)" for (base_model, init_model, x_size) in zip(base_models,
             init_models, x_sizes)
-            @info solver, mtype, jacobian_regularization, base_model, init_model, x_size
-
             model = if mtype === :deq
                 DeepEquilibriumNetwork(base_model, solver; jacobian_regularization)
             elseif mtype === :skipdeq
@@ -48,9 +48,8 @@ end
             x = randn(rng, Float32, x_size...) |> dev
             z, st = model(x, ps, st)
 
-            opt_broken = solver isa NewtonRaphson ||
-                         solver isa SimpleLimitedMemoryBroyden
-            @jet model(x, ps, st) opt_broken=opt_broken # Broken due to nfe dynamic dispatch
+            opt_broken = solver isa SimpleLimitedMemoryBroyden
+            @jet model(x, ps, st) opt_broken=opt_broken
 
             @test all(isfinite, z)
             @test size(z) == size(x)
@@ -107,20 +106,18 @@ end
     scales = [((4,), (3,), (2,), (1,))]
 
     model_type = (:deq, :skipdeq, :skipregdeq, :node)
-    solvers = (VCAB3(), Tsit5(), NewtonRaphson(), SimpleLimitedMemoryBroyden())
+    solvers = (VCAB3(), Tsit5(),
+        NewtonRaphson(; autodiff=AutoForwardDiff(; chunksize=12)),
+        SimpleLimitedMemoryBroyden())
     jacobian_regularizations = (nothing,)
 
     for mtype in model_type, jacobian_regularization in jacobian_regularizations
         @testset "Solver: $(__nameof(solver))" for solver in solvers
             @testset "x_size: $(x_size)" for (main_layer, mapping_layer, init_layer, x_size, scale) in zip(main_layers,
                 mapping_layers, init_layers, x_sizes, scales)
-                @info solver, mtype, jacobian_regularization, main_layer, mapping_layer,
-                init_layer, x_size, scale
-
                 model = if mtype === :deq
                     MultiScaleDeepEquilibriumNetwork(main_layer, mapping_layer, nothing,
-                        solver,
-                        scale; jacobian_regularization)
+                        solver, scale; jacobian_regularization)
                 elseif mtype === :skipdeq
                     MultiScaleSkipDeepEquilibriumNetwork(main_layer, mapping_layer, nothing,
                         init_layer, solver, scale; jacobian_regularization)
@@ -140,8 +137,7 @@ end
                 z, st = model(x, ps, st)
                 z_ = DEQs.__flatten_vcat(z)
 
-                opt_broken = solver isa NewtonRaphson ||
-                             solver isa SimpleLimitedMemoryBroyden
+                opt_broken = solver isa SimpleLimitedMemoryBroyden
                 @jet model(x, ps, st) opt_broken=opt_broken # Broken due to nfe dynamic dispatch
 
                 @test all(isfinite, z_)

test/qa.jl

@@ -3,6 +3,5 @@ import ChainRulesCore as CRC
 
 @testset "Aqua" begin
     Aqua.test_all(DeepEquilibriumNetworks; ambiguities=false)
-    Aqua.test_ambiguities(DeepEquilibriumNetworks; recursive=false,
-        exclude=[CRC.rrule, CRC.frule])
+    Aqua.test_ambiguities(DeepEquilibriumNetworks; recursive=false)
 end

test/runtests.jl

@@ -1,13 +1,7 @@
-using SafeTestsets, Test
+using SafeTestsets, Test, TestSetExtensions
 
-@testset "Deep Equilibrium Networks" begin
-    @safetestset "Quality Assurance" begin
-        include("qa.jl")
-    end
-    @safetestset "Utilities" begin
-        include("utils.jl")
-    end
-    @safetestset "Layers" begin
-        include("layers.jl")
-    end
+@testset ExtendedTestSet "Deep Equilibrium Networks" begin
+    @safetestset "Quality Assurance" include("qa.jl")
+    @safetestset "Utilities" include("utils.jl")
+    @safetestset "Layers" include("layers.jl")
 end