Cannot simulate single input `SymbolicNeuralNetwork` model

[TorkelE]

@SebastianM-C

Unfortunatley I cannot get the new symbolic interface to work. This example, adapted from the tests, work:

# Fetch packages.
using ModelingToolkit, ModelingToolkitNeuralNets
using ModelingToolkit: t_nounits as t, D_nounits as D
using OrdinaryDiffEq
using StableRNGs

let
    # Make Neural network.
    chain = multi_layer_feed_forward(2, 2)
    NN, p = SymbolicNeuralNetwork(; chain, n_input = 2, n_output = 2, rng = StableRNG(42))

    # Create UDE model.
    @variables x(t) y(t)
    @parameters α δ
    eqs = [
        D(x) ~ α * x + NN([x, y], p)[1],
        D(y) ~ -δ * y + NN([x, y], p)[2]]
    @mtkbuild osys = ODESystem(eqs, ModelingToolkit.t_nounits)

    # Test UDE simulation.
    prob = ODEProblem(osys, [x => 3.1, y => 1.5], (0, 1.0), [α => 1.3 δ => 1.8])
    sol = solve(prob)
end

However, if I try to create a model with only a single input to the neural network (here changing n_input to 1, and making similar modification in the call) I get an error:

let
    # Make Neural network.
    chain = multi_layer_feed_forward(2, 2)
    NN, p = SymbolicNeuralNetwork(; chain, n_input = 1, n_output = 2, rng = StableRNG(42))

    # Create UDE model.
    @variables x(t) y(t)
    @parameters α δ
    eqs = [
        D(x) ~ α * x + NN([x], p)[1],
        D(y) ~ -δ * y + NN([x], p)[2]]
    @mtkbuild osys = ODESystem(eqs, ModelingToolkit.t_nounits)

    # Test UDE simulation.
    prob = ODEProblem(osys, [x => 3.1, y => 1.5], (0, 1.0), [α => 1.3 δ => 1.8])
    sol = solve(prob)
end

ERROR: DimensionMismatch: A has dimensions (5,2) but B has dimensions (1,1)
Stacktrace:
  [1] gemm_wrapper!(C::Matrix{…}, tA::Char, tB::Char, A::Base.ReshapedArray{…}, B::Matrix{…}, _add::LinearAlgebra.MulAddMul{…})
    @ LinearAlgebra ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:629
  [2] generic_matmatmul!
    @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:381 [inlined]
  [3] _mul!
    @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:287 [inlined]
  [4] mul!
    @ ~/.julia/juliaup/julia-1.11.5+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/matmul.jl:285 [inlined]
  [5] matmul_linalg_default!
    @ ~/.julia/packages/LuxLib/EnsP3/src/impl/matmul.jl:205 [inlined]
  [6] matmul_cpu_fallback!
    @ ~/.julia/packages/LuxLib/EnsP3/src/impl/matmul.jl:177 [inlined]
  [7] matmul_cpu!
    @ ~/.julia/packages/LuxLib/EnsP3/src/impl/matmul.jl:164 [inlined]
  [8] matmul!
    @ ~/.julia/packages/LuxLib/EnsP3/src/impl/matmul.jl:124 [inlined]
  [9] fused_dense!
    @ ~/.julia/packages/LuxLib/EnsP3/src/impl/dense.jl:54 [inlined]
 [10] fused_dense
    @ ~/.julia/packages/LuxLib/EnsP3/src/impl/dense.jl:42 [inlined]
 [11] fused_dense
    @ ~/.julia/packages/LuxLib/EnsP3/src/impl/dense.jl:16 [inlined]
 [12] fused_dense_bias_activation
    @ ~/.julia/packages/LuxLib/EnsP3/src/api/dense.jl:36 [inlined]
 [13] (::Lux.Dense{…})(x::Vector{…}, ps::ComponentArrays.ComponentVector{…}, st::@NamedTuple{})
    @ Lux ~/.julia/packages/Lux/CYnn3/src/layers/basic.jl:357
 [14] apply
    @ ~/.julia/packages/LuxCore/Av7WJ/src/LuxCore.jl:155 [inlined]
 [15] macro expansion
    @ ~/.julia/packages/Lux/CYnn3/src/layers/containers.jl:0 [inlined]
 [16] applychain
    @ ~/.julia/packages/Lux/CYnn3/src/layers/containers.jl:511 [inlined]
 [17] Chain
    @ ~/.julia/packages/Lux/CYnn3/src/layers/containers.jl:509 [inlined]
 [18] apply
    @ ~/.julia/packages/LuxCore/Av7WJ/src/LuxCore.jl:155 [inlined]
 [19] stateless_apply(model::Lux.Chain{…}, x::Vector{…}, ps::ComponentArrays.ComponentVector{…})
    @ LuxCore ~/.julia/packages/LuxCore/Av7WJ/src/LuxCore.jl:166
 [20] (::ModelingToolkitNeuralNets.StatelessApplyWrapper{…})(input::Vector{…}, nn_p::SubArray{…})
    @ ModelingToolkitNeuralNets ~/.julia/packages/ModelingToolkitNeuralNets/qmb5j/src/ModelingToolkitNeuralNets.jl:120
 [21] macro expansion
    @ ~/.julia/packages/SymbolicUtils/NGWJM/src/code.jl:510 [inlined]
 [22] macro expansion
    @ ~/.julia/packages/Symbolics/B6Z8m/src/build_function.jl:368 [inlined]
 [23] macro expansion
    @ ~/.julia/packages/RuntimeGeneratedFunctions/M9ZX8/src/RuntimeGeneratedFunctions.jl:163 [inlined]
 [24] macro expansion
    @ ./none:0 [inlined]
 [25] generated_callfunc
    @ ./none:0 [inlined]
 [26] (::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{…})(::Vector{…}, ::Vector{…}, ::MTKParameters{…}, ::Float64)
    @ RuntimeGeneratedFunctions ~/.julia/packages/RuntimeGeneratedFunctions/M9ZX8/src/RuntimeGeneratedFunctions.jl:150
 [27] macro expansion
    @ ~/.julia/packages/ModelingToolkit/aau6A/src/systems/codegen_utils.jl:0 [inlined]
 [28] _generated_call
    @ ~/.julia/packages/ModelingToolkit/aau6A/src/systems/codegen_utils.jl:262 [inlined]
 [29] GeneratedFunctionWrapper
    @ ~/.julia/packages/ModelingToolkit/aau6A/src/systems/codegen_utils.jl:259 [inlined]
 [30] Void
    @ ~/.julia/packages/SciMLBase/c6Noy/src/utils.jl:486 [inlined]
 [31] (::FunctionWrappers.CallWrapper{…})(f::SciMLBase.Void{…}, arg1::Vector{…}, arg2::Vector{…}, arg3::MTKParameters{…}, arg4::Float64)
    @ FunctionWrappers ~/.julia/packages/FunctionWrappers/Q5cBx/src/FunctionWrappers.jl:65
 [32] macro expansion
    @ ~/.julia/packages/FunctionWrappers/Q5cBx/src/FunctionWrappers.jl:137 [inlined]
 [33] do_ccall
    @ ~/.julia/packages/FunctionWrappers/Q5cBx/src/FunctionWrappers.jl:125 [inlined]
 [34] FunctionWrapper
    @ ~/.julia/packages/FunctionWrappers/Q5cBx/src/FunctionWrappers.jl:144 [inlined]
 [35] _call
    @ ~/.julia/packages/FunctionWrappersWrappers/9XR0m/src/FunctionWrappersWrappers.jl:12 [inlined]
 [36] FunctionWrappersWrapper
    @ ~/.julia/packages/FunctionWrappersWrappers/9XR0m/src/FunctionWrappersWrappers.jl:10 [inlined]
--- the above 7 lines are repeated 1 more time ---
 [44] ODEFunction
    @ ~/.julia/packages/SciMLBase/c6Noy/src/scimlfunctions.jl:2470 [inlined]
 [45] initialize!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…}, cache::OrdinaryDiffEqTsit5.Tsit5Cache{…})
    @ OrdinaryDiffEqTsit5 ~/.julia/packages/OrdinaryDiffEqTsit5/DHYtz/src/tsit_perform_step.jl:175
 [46] initialize!(integrator::OrdinaryDiffEqCore.ODEIntegrator{…}, cache::OrdinaryDiffEqCore.DefaultCache{…})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/bMOsj/src/perform_step/composite_perform_step.jl:38
 [47] __init(prob::ODEProblem{…}, alg::CompositeAlgorithm{…}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{…}; saveat::Tuple{}, tstops::Tuple{}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::Nothing, dense::Bool, calck::Bool, dt::Float64, dtmin::Float64, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Rational{…}, abstol::Nothing, reltol::Nothing, qmin::Rational{…}, qmax::Int64, qsteady_min::Int64, qsteady_max::Int64, beta1::Nothing, beta2::Nothing, qoldinit::Rational{…}, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(LinearAlgebra.opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), progress_id::Symbol, userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias::ODEAliasSpecifier, initializealg::OrdinaryDiffEqCore.DefaultInit, kwargs::@Kwargs{…})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/bMOsj/src/solve.jl:587
 [48] __init (repeats 5 times)
    @ ~/.julia/packages/OrdinaryDiffEqCore/bMOsj/src/solve.jl:11 [inlined]
 [49] #__solve#62
    @ ~/.julia/packages/OrdinaryDiffEqCore/bMOsj/src/solve.jl:6 [inlined]
 [50] __solve
    @ ~/.julia/packages/OrdinaryDiffEqCore/bMOsj/src/solve.jl:1 [inlined]
 [51] solve_call(_prob::ODEProblem{…}, args::CompositeAlgorithm{…}; merge_callbacks::Bool, kwargshandle::Nothing, kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:635
 [52] solve_up(prob::ODEProblem{…}, sensealg::Nothing, u0::Vector{…}, p::MTKParameters{…}, args::CompositeAlgorithm{…}; kwargs::@Kwargs{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1142
 [53] solve_up
    @ ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1120 [inlined]
 [54] #solve#42
    @ ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1059 [inlined]
 [55] solve
    @ ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1047 [inlined]
 [56] #__solve#3
    @ ~/.julia/packages/OrdinaryDiffEqDefault/kOV75/src/default_alg.jl:48 [inlined]
 [57] __solve
    @ ~/.julia/packages/OrdinaryDiffEqDefault/kOV75/src/default_alg.jl:47 [inlined]
 [58] #__solve#63
    @ ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1451 [inlined]
 [59] __solve
    @ ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1442 [inlined]
 [60] #solve_call#35
    @ ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:635 [inlined]
 [61] solve_call(::ODEProblem{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:592
 [62] solve_up(::ODEProblem{…}, ::Nothing, ::Vector{…}, ::MTKParameters{…}; kwargs::@Kwargs{})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1126
 [63] solve_up
    @ ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1120 [inlined]
 [64] solve(::ODEProblem{…}; sensealg::Nothing, u0::Nothing, p::Nothing, wrap::Val{…}, kwargs::@Kwargs{})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1057
 [65] solve(::ODEProblem{…})
    @ DiffEqBase ~/.julia/packages/DiffEqBase/zYZst/src/solve.jl:1047
Some type information was truncated. Use `show(err)` to see complete types.

(I originally I had a different model where I tried to get SymbolicNeuralNetwork to work, but that failed similarly, but started working of I increased n_input to 2).

[ChrisRackauckas]

Symbolic array stuff might be doing something funky for size=1. That's worth digging into.

[SebastianM-C]

Note that

chain = multi_layer_feed_forward(2, 2)

has to match n_input = 1, n_output = 2.

I was thinking if we can make it less error prone by trying to get the sizes from the Lux model, but only the ouputsize is defined generically. Input size is trickier to define and I think it was rejected in Lux when I tried to add something for that.

[SebastianM-C]

Let me know if you also have this issue if you change to multi_layer_feed_forward(1, 2). That should fix it.

[TorkelE]

Ahh, that makes sense! I thought that multi_layer_feed_forward was the internal, and then n_input and n_output was apended to that chain.

Is there any reason that n_input and n_output cannot be infered from the chain and cannot be gved as well?

[SebastianM-C]

Is there any reason that n_input and n_output cannot be infered from the chain and cannot be gved as well?

The reason is that we need LuxCore interfaces for that and you can't define the input size just from the information of the chain. The output can be defined more generically because you pass something through the network and see what comes out.

For Chain and Dense this is trivial, but to have the interface in LuxCore, it would need to apply to more cases.

[TorkelE]

Ok. For the limited cases of Chain, is there an easy way to retrieve these parts from the chain which I can insert into my modelling pipeline for now to handle this (and potentially also use in relevant tutorials)?

[TorkelE]

(also, thanks for the ultra quick reply!)

[SebastianM-C]

One compromise would be for us to try to get this info in SymbolicNeuralNetwork / NeuralNetworkBlock for a limited set of cases where we can guarantee the correct answer (Like for Chain and Dense). Or we could maybe try to error earlier if we detect a mismatch 🤔

[TorkelE]

Both would definitely be useful!
(although also partially easier for users when we have proper docs and tutorials ready)

[SebastianM-C]

As mentioned in https://lux.csail.mit.edu/stable/api/Building_Blocks/LuxCore , outputsize(layer) was intentionally removed in the 1.0 release of Lux (and inputsize(layer) didn't even make it). In the constructors we could try doing a outputsize(layer, x, rng) call with x having the input size specified by the user and this can either error like in your case or give the output size that can be compared with the n_output.

I'll see what I can do 🤔

[TorkelE]

Sounds good, and that could definitely be helpful. I got things to work now, but I will probably have more queries as I play around more with the new interface