diff --git a/docs/src/friction.md b/docs/src/friction.md index e9a6f21..c59beb4 100644 --- a/docs/src/friction.md +++ b/docs/src/friction.md @@ -96,7 +96,7 @@ chain = Lux.Chain( Lux.Dense(10 => 10, Lux.mish, use_bias = false), Lux.Dense(10 => 1, use_bias = false) ) -nn = NeuralNetworkBlock(1, 1; chain = chain, rng = StableRNG(1111)) +@named nn = NeuralNetworkBlock(1, 1; chain = chain, rng = StableRNG(1111)) eqs = [connect(model.nn_in, nn.output) connect(model.nn_out, nn.input)] diff --git a/src/ModelingToolkitNeuralNets.jl b/src/ModelingToolkitNeuralNets.jl index fd55d09..b8c9c17 100644 --- a/src/ModelingToolkitNeuralNets.jl +++ b/src/ModelingToolkitNeuralNets.jl @@ -17,7 +17,8 @@ include("utils.jl") chain = multi_layer_feed_forward(n_input, n_output), rng = Xoshiro(0), init_params = Lux.initialparameters(rng, chain), - eltype = Float64) + eltype = Float64, + name) Create an `ODESystem` with a neural network inside. """ @@ -26,7 +27,8 @@ function NeuralNetworkBlock(n_input = 1, chain = multi_layer_feed_forward(n_input, n_output), rng = Xoshiro(0), init_params = Lux.initialparameters(rng, chain), - eltype = Float64) + eltype = Float64, + name) ca = ComponentArray{eltype}(init_params) @parameters p[1:length(ca)] = Vector(ca) @@ -39,8 +41,8 @@ function NeuralNetworkBlock(n_input = 1, eqs = [output.u ~ out] - @named ude_comp = ODESystem( - eqs, t_nounits, [], [p, T], systems = [input, output]) + ude_comp = ODESystem( + eqs, t_nounits, [], [p, T]; systems = [input, output], name) return ude_comp end diff --git a/test/lotka_volterra.jl b/test/lotka_volterra.jl index 6b6f7a9..51ea9ef 100644 --- a/test/lotka_volterra.jl +++ b/test/lotka_volterra.jl @@ -44,7 +44,7 @@ end model = lotka_ude() chain = multi_layer_feed_forward(2, 2) -nn = NeuralNetworkBlock(2, 2; chain, rng = StableRNG(42)) +@named nn = NeuralNetworkBlock(2, 2; chain, rng = StableRNG(42)) eqs = [connect(model.nn_in, nn.output) connect(model.nn_out, nn.input)]