Add PrecompileTools workload

Project.toml

@@ -9,6 +9,7 @@ IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLPublic = "431bcebd-1456-4ced-9d72-93c2757fff0b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
@@ -30,6 +31,7 @@ OptimizationBase = "5"
 OptimizationOptimJL = "0.4.13"
 OptimizationOptimisers = "0.3.14"
 OrdinaryDiffEqVerner = "2.3.0"
+PrecompileTools = "1.2.1"
 Random = "1.10"
 SafeTestsets = "0.1"
 SciCompDSL = "1"

src/ModelingToolkitNeuralNets.jl

@@ -9,6 +9,7 @@ using LuxCore: initialparameters, stateless_apply, outputsize
 using Lux: Lux
 using Random: Xoshiro
 using ComponentArrays: ComponentArray
+using PrecompileTools: @compile_workload, @setup_workload
 using SciMLPublic: @public
 
 export NeuralNetworkBlock, SymbolicNeuralNetwork, @SymbolicNeuralNetwork, multi_layer_feed_forward, get_network
@@ -314,6 +315,8 @@ _num_chain_outputs(chain::Lux.Chain) = _num_layer_outputs(chain.layers[end])
 _num_layer_outputs(layer::Lux.Dense) = layer.out_dims
 _num_layer_outputs(layer) = error("@SymbolicNeuralNetwork has been provided with a chain which last layer's type ($(typeof(layer))) is not supported for automatic output size detection. Please use the `SymbolicNeuralNetwork` function directly.")
 
+include("precompile.jl")
+
 # Layer types that can potentially be supported in the future.
 # _num_layer_inputs(layer::Lux.Bilinear) = layer.in1_dims + layer.in2_dims
 # _num_layer_inputs(layer::Lux.RNNCell) = layer.in_dims

src/precompile.jl

@@ -0,0 +1,9 @@
+@setup_workload begin
+    chain = multi_layer_feed_forward(2, 1; width = 3, depth = 1, activation = tanh)
+    rng = Xoshiro(0)
+
+    @compile_workload begin
+        SymbolicNeuralNetwork(; chain, n_input = 2, n_output = 1, rng)
+        NeuralNetworkBlock(2, 1; chain, rng, name = :nn)
+    end
+end

test/precompile_workload.jl

@@ -0,0 +1,13 @@
+using Lux
+using ModelingToolkitBase
+using ModelingToolkitNeuralNets
+using Random
+using Test
+
+chain = multi_layer_feed_forward(2, 1; width = 3, depth = 1, activation = tanh)
+NN, p = SymbolicNeuralNetwork(; chain, n_input = 2, n_output = 1, rng = Xoshiro(0))
+nn = NeuralNetworkBlock(2, 1; chain, rng = Xoshiro(0), name = :nn)
+
+@test get_network(ModelingToolkitBase.getdefault(NN)) === chain
+@test length(p) > 0
+@test ModelingToolkitBase.getname(nn) == :nn

test/runtests.jl

@@ -3,6 +3,7 @@ using SciMLTesting
 
 run_tests(;
     core = function ()
+        @safetestset "Precompile workload" include("precompile_workload.jl")
         @safetestset "Basic" include("lotka_volterra.jl")
         @safetestset "MTK model macro compatibility" include("macro.jl")
         @safetestset "Symbolic Neural Network Macro" include("symbolicnn_macro.jl")