Drop the jacreg for multi scale models

Author: avik-pal

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0-rc2"
 manifest_format = "2.0"
-project_hash = "014ca85a3845de2d1ab4491eb2d8b2f946e37a20"
+project_hash = "8d33856ceafbd0190b2f9c40c1bcfddbb51890fa"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "332e5d7baeff8497b923b730b994fa480601efc7"
@@ -696,9 +696,11 @@ uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
 [[deps.LinearSolve]]
 deps = ["ArrayInterface", "ConcreteStructs", "DocStringExtensions", "EnumX", "FastLapackInterface", "GPUArraysCore", "InteractiveUtils", "KLU", "Krylov", "Libdl", "LinearAlgebra", "MKL_jll", "PrecompileTools", "Preferences", "RecursiveFactorization", "Reexport", "Requires", "SciMLBase", "SciMLOperators", "Setfield", "SparseArrays", "Sparspak", "StaticArraysCore", "UnPack"]
-git-tree-sha1 = "03cd36673f050ec49f9bdfce31e0e7e22d877e64"
+git-tree-sha1 = "7c9f62dc6f7d11b0f8dfb3c1f3637cb220dc9866"
+repo-rev = "ap/normal_cholesky_dispatches"
+repo-url = "https://github.com/SciML/LinearSolve.jl.git"
 uuid = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
-version = "2.21.0"
+version = "2.21.1"
 
     [deps.LinearSolve.extensions]
     LinearSolveBandedMatricesExt = "BandedMatrices"

Project.toml

@@ -10,6 +10,7 @@ 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"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"

src/layers.jl

@@ -306,7 +306,9 @@ julia> model(x, ps, st);
 ```
 """
 function MultiScaleDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers::Matrix,
-        post_fuse_layer::Union{Nothing, Tuple}, solver, scales; kwargs...)
+        post_fuse_layer::Union{Nothing, Tuple}, solver, scales;
+        jacobian_regularization=nothing, kwargs...)
+    @assert jacobian_regularization === nothing "Jacobian Regularization is not supported yet for MultiScale Models."
     l1 = Parallel(nothing, main_layers...)
     l2 = BranchLayer(Parallel.(+, map(x -> tuple(x...), eachrow(mapping_layers))...)...)
 

test/layers.jl

@@ -108,7 +108,7 @@ end
 
     model_type = (:deq, :skipdeq, :skipregdeq, :node)
     solvers = (VCAB3(), Tsit5(), NewtonRaphson(), SimpleLimitedMemoryBroyden())
-    jacobian_regularizations = (nothing, AutoFiniteDiff(), AutoZygote())
+    jacobian_regularizations = (nothing,)
 
     for mtype in model_type, jacobian_regularization in jacobian_regularizations
         @testset "Solver: $(__nameof(solver))" for solver in solvers