Update to MTK@v10 and drop the dependency on MTKStdLib

We can simplify the NeuralNetworkBlock a bit by using normal vector
variables instead of RealInput/OutputArrays.

Author: SebastianM-C

Project.toml

@@ -1,41 +1,41 @@
 name = "ModelingToolkitNeuralNets"
 uuid = "f162e290-f571-43a6-83d9-22ecc16da15f"
 authors = ["Sebastian Micluța-Câmpeanu <[email protected]> and contributors"]
-version = "1.7.0"
+version = "2.0.0-DEV"
 
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
-ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
 Aqua = "0.8"
 ComponentArrays = "0.15.11"
-DifferentiationInterface = "0.6"
-ForwardDiff = "0.10.36"
+DifferentiationInterface = "0.6, 0.7"
+ForwardDiff = "0.10.36, 1"
 IntervalSets = "0.7.10"
-JET = "0.8, 0.9"
-Lux = "1"
-LuxCore = "1"
-ModelingToolkit = "9.64"
-ModelingToolkitStandardLibrary = "2.7"
-Optimization = "3.24, 4"
+JET = "0.8, 0.9, 0.10"
+Lux = "1.14"
+LuxCore = "1.2"
+ModelingToolkit = "10"
+ModelingToolkitStandardLibrary = "2.24"
+Optimization = "4"
 OptimizationOptimisers = "0.2.1, 0.3"
-OrdinaryDiffEq = "6.74"
+OrdinaryDiffEqVerner = "1.2"
 Random = "1.10"
 SafeTestsets = "0.1"
 SciMLSensitivity = "7.72"
 SciMLStructures = "1.1.0"
 StableRNGs = "1"
-SymbolicIndexingInterface = "0.3.15"
-Symbolics = "6.36"
+Statistics = "1.10"
+SymbolicIndexingInterface = "0.3.41"
+Symbolics = "6.43"
 Test = "1.10"
-Zygote = "0.6.73"
+Zygote = "0.6.73, 0.7"
 julia = "1.10"
 
 [extras]
@@ -45,14 +45,19 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimisers = "42dfb2eb-d2b4-4451-abcd-913932933ac1"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqVerner = "79d7bb75-1356-48c1-b8c0-6832512096c2"
+ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "JET", "Test", "OrdinaryDiffEq", "DifferentiationInterface", "SciMLSensitivity", "Zygote", "ForwardDiff", "Optimization", "OptimizationOptimisers", "SafeTestsets", "SciMLStructures", "StableRNGs", "SymbolicIndexingInterface"]
+test = ["Aqua", "JET", "Test", "OrdinaryDiffEqVerner", "DifferentiationInterface",
+    "SciMLSensitivity", "Zygote", "ForwardDiff", "ModelingToolkitStandardLibrary",
+    "Optimization", "OptimizationOptimisers", "SafeTestsets", "SciMLStructures",
+    "StableRNGs", "Statistics", "SymbolicIndexingInterface"]

src/ModelingToolkitNeuralNets.jl

@@ -1,10 +1,9 @@
 module ModelingToolkitNeuralNets
 
-using ModelingToolkit: @parameters, @named, ODESystem, t_nounits
+using ModelingToolkit: @parameters, @named, @variables, System, t_nounits
 using IntervalSets: var".."
-using ModelingToolkitStandardLibrary.Blocks: RealInputArray, RealOutputArray
 using Symbolics: Symbolics, @register_array_symbolic, @wrapped
-using LuxCore: stateless_apply
+using LuxCore: stateless_apply, outputsize
 using Lux: Lux
 using Random: Xoshiro
 using ComponentArrays: ComponentArray
@@ -21,7 +20,7 @@ include("utils.jl")
         eltype = Float64,
         name)
 
-Create an `ODESystem` with a neural network inside.
+Create a component neural network as a `System`.
 """
 function NeuralNetworkBlock(; n_input = 1, n_output = 1,
         chain = multi_layer_feed_forward(n_input, n_output),
@@ -31,19 +30,21 @@ function NeuralNetworkBlock(; n_input = 1, n_output = 1,
         name)
     ca = ComponentArray{eltype}(init_params)
 
-    @parameters p[1:length(ca)] = Vector(ca)
+    @parameters p[1:length(ca)]=Vector(ca) [tunable = true]
     @parameters T::typeof(typeof(ca))=typeof(ca) [tunable = false]
-    @parameters lux_model::typeof(chain) = chain
+    @parameters lux_model::typeof(chain)=chain [tunable = false]
 
-    @named input = RealInputArray(nin = n_input)
-    @named output = RealOutputArray(nout = n_output)
+    @variables inputs(t_nounits)[1:n_input] [input = true]
+    @variables outputs(t_nounits)[1:n_output] [output = true]
 
-    out = stateless_apply(lux_model, input.u, lazyconvert(T, p))
+    expected_outsz = only(outputsize(chain, inputs, rng))
+    msg = "The outputsize of the given Lux network ($expected_outsz) does not match `n_output = $n_output`"
+    @assert n_output==expected_outsz msg
 
-    eqs = [output.u ~ out]
+    eqs = [outputs ~ stateless_apply(lux_model, inputs, lazyconvert(T, p))]
 
-    ude_comp = ODESystem(
-        eqs, t_nounits, [], [lux_model, p, T]; systems = [input, output], name)
+    ude_comp = System(
+        eqs, t_nounits, [inputs, outputs], [lux_model, p, T]; name)
     return ude_comp
 end