refactor: rename create_ude_component to NeuralNetworkBlock

Author: SebastianM-C

src/UDEComponents.jl

@@ -8,23 +8,23 @@ using Lux: Lux
 using Random: Xoshiro
 using ComponentArrays: ComponentArray
 
-export create_ude_component, multi_layer_feed_forward
+export NeuralNetworkBlock, multi_layer_feed_forward
 
 include("utils.jl")
 
 """
-    create_ude_component(n_input = 1, n_output = 1;
+    NeuralNetworkBlock(n_input = 1, n_output = 1;
         chain = multi_layer_feed_forward(n_input, n_output),
         rng = Xoshiro(0))
 
 Create an `ODESystem` with a neural network inside.
 """
-function create_ude_component(n_input = 1,
+function NeuralNetworkBlock(n_input = 1,
         n_output = 1;
         chain = multi_layer_feed_forward(n_input, n_output),
-        rng = Xoshiro(0))
-    lux_p, st = Lux.setup(rng, chain)
-    ca = ComponentArray(lux_p)
+        rng = Xoshiro(0), eltype=Float64)
+    lux_p = Lux.initialparameters(rng, chain)
+    ca = ComponentArray{eltype}(lux_p)
 
     @parameters p[1:length(ca)] = Vector(ca)
     @parameters T::typeof(typeof(p))=typeof(p) [tunable = false]