Neural ODEs on GPUs

Note that the differential equation solvers will run on the GPU if the initial condition is a GPU array. Thus, for example, we can define a neural ODE manually that runs on the GPU (if no GPU is available, the calculation defaults back to the CPU).

For a detailed discussion on how GPUs need to be setup refer to Lux Docs.

using OrdinaryDiffEq, Lux, LuxCUDA, SciMLSensitivity, ComponentArrays, Random
rng = Xoshiro(0)

const cdev = cpu_device()
const gdev = gpu_device()

model = Chain(Dense(2, 50, tanh), Dense(50, 2))
ps, st = Lux.setup(rng, model)
ps = ps |> ComponentArray |> gdev
st = st |> gdev
dudt(u, p, t) = model(u, p, st)[1]

# Simulation interval and intermediary points
tspan = (0.0f0, 10.0f0)
tsteps = 0.0f0:1.0f-1:10.0f0

u0 = Float32[2.0; 0.0] |> gdev
prob_gpu = ODEProblem(dudt, u0, tspan, ps)

# Runs on a GPU
sol_gpu = solve(prob_gpu, Tsit5(); saveat = tsteps)

retcode: Success
Interpolation: 1st order linear
t: 101-element Vector{Float32}:
  0.0
  0.1
  0.2
  0.3
  0.4
  0.5
  0.6
  0.7
  0.8
  0.9
  ⋮
  9.2
  9.3
  9.4
  9.5
  9.6
  9.7
  9.8
  9.9
 10.0
u: 101-element Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}:
 Float32[2.0, 0.0]
 Float32[1.7873143, 0.031452507]
 Float32[1.5787975, 0.06454538]
 Float32[1.3766681, 0.09886337]
 Float32[1.1836983, 0.13350704]
 Float32[1.0028386, 0.16698027]
 Float32[0.8367471, 0.19752759]
 Float32[0.6875932, 0.22379528]
 Float32[0.5568672, 0.24555919]
 Float32[0.445041, 0.26368508]
 ⋮
 Float32[1.8026602, 12.04514]
 Float32[1.8133537, 12.216526]
 Float32[1.8231647, 12.387732]
 Float32[1.8320737, 12.558775]
 Float32[1.8400607, 12.729679]
 Float32[1.8471066, 12.900487]
 Float32[1.8531883, 13.071189]
 Float32[1.8582877, 13.241817]
 Float32[1.8623875, 13.412399]

Or we could directly use the neural ODE layer function, like:

using DiffEqFlux: NeuralODE
prob_neuralode_gpu = NeuralODE(model, tspan, Tsit5(); saveat = tsteps)

NeuralODE(
    model = Chain(
        layer_1 = Dense(2 => 50, tanh),  # 150 parameters
        layer_2 = Dense(50 => 2),       # 102 parameters
    ),
)         # Total: 252 parameters,
          #        plus 0 states.

If one is using Lux.Chain, then the computation takes place on the GPU with f(x,p,st) if x, p and st are on the GPU. This commonly looks like:

dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))

u0 = Float32[2.0; 0.0] |> gdev
p, st = Lux.setup(rng, dudt2) |> gdev

dudt2_(u, p, t) = first(dudt2(u, p, st))

# Simulation interval and intermediary points
tspan = (0.0f0, 10.0f0)
tsteps = 0.0f0:1.0f-1:10.0f0

prob_gpu = ODEProblem(dudt2_, u0, tspan, p)

# Runs on a GPU
sol_gpu = solve(prob_gpu, Tsit5(); saveat = tsteps)

retcode: Success
Interpolation: 1st order linear
t: 101-element Vector{Float32}:
  0.0
  0.1
  0.2
  0.3
  0.4
  0.5
  0.6
  0.7
  0.8
  0.9
  ⋮
  9.2
  9.3
  9.4
  9.5
  9.6
  9.7
  9.8
  9.9
 10.0
u: 101-element Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}:
 Float32[2.0, 0.0]
 Float32[2.0003984, 0.084881306]
 Float32[2.000786, 0.16975331]
 Float32[2.0011172, 0.2546237]
 Float32[2.0013137, 0.33950973]
 Float32[2.0012627, 0.42445177]
 Float32[2.0008066, 0.5095357]
 Float32[1.9997556, 0.5948871]
 Float32[1.9978168, 0.6808015]
 Float32[1.9945722, 0.76778114]
 ⋮
 Float32[-0.44770268, 0.3272455]
 Float32[-0.46191934, 0.35120934]
 Float32[-0.47744444, 0.37346473]
 Float32[-0.49423942, 0.39385107]
 Float32[-0.51219267, 0.41223925]
 Float32[-0.5311202, 0.42853218]
 Float32[-0.5507664, 0.44266585]
 Float32[-0.570801, 0.4546066]
 Float32[-0.5908468, 0.46434626]

or via the NeuralODE struct:

prob_neuralode_gpu = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)
prob_neuralode_gpu(u0, p, st)

(SciMLBase.ODESolution{Float32, 2, Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}, Nothing, Nothing, Vector{Float32}, Vector{Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}}, Nothing, SciMLBase.ODEProblem{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}, Tuple{Float32, Float32}, false, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{weight::CUDA.CuArray{Float32, 2, CUDA.DeviceMemory}, bias::CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}, layer_3::@NamedTuple{weight::CUDA.CuArray{Float32, 2, CUDA.DeviceMemory}, bias::CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}}, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#dudt#17"{StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}, OrdinaryDiffEqTsit5.Tsit5{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), Static.False}, OrdinaryDiffEqCore.InterpolationData{SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#dudt#17"{StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}, Vector{Float32}, Vector{Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}}, Nothing, OrdinaryDiffEqTsit5.Tsit5ConstantCache, Nothing}, SciMLBase.DEStats, Nothing, Nothing, Nothing, Nothing}(CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}[Float32[2.0, 0.0], Float32[2.0003984, 0.084881306], Float32[2.000786, 0.16975331], Float32[2.0011172, 0.2546237], Float32[2.0013137, 0.33950973], Float32[2.0012627, 0.42445177], Float32[2.0008066, 0.5095357], Float32[1.9997556, 0.5948871], Float32[1.9978168, 0.6808015], Float32[1.9945722, 0.76778114]  …  Float32[-0.4347611, 0.30176672], Float32[-0.44770268, 0.3272455], Float32[-0.46191934, 0.35120934], Float32[-0.47744444, 0.37346473], Float32[-0.49423942, 0.39385107], Float32[-0.51219267, 0.41223925], Float32[-0.5311202, 0.42853218], Float32[-0.5507664, 0.44266585], Float32[-0.570801, 0.4546066], Float32[-0.5908468, 0.46434626]], nothing, nothing, Float32[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9  …  9.1, 9.2, 9.3, 9.4, 9.5, 9.6, 9.7, 9.8, 9.9, 10.0], Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}[[Float32[2.0, 0.0]]], nothing, SciMLBase.ODEProblem{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}, Tuple{Float32, Float32}, false, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{weight::CUDA.CuArray{Float32, 2, CUDA.DeviceMemory}, bias::CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}, layer_3::@NamedTuple{weight::CUDA.CuArray{Float32, 2, CUDA.DeviceMemory}, bias::CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}}, SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#dudt#17"{StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, SciMLBase.StandardODEProblem}(SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#dudt#17"{StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(DiffEqFlux.var"#dudt#17"{StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}(StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}(Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = WrappedFunction(#1), layer_2 = Dense(2 => 50, tanh), layer_3 = Dense(50 => 2)), nothing), nothing, (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple()), nothing, static(true))), LinearAlgebra.UniformScaling{Bool}(true), nothing, DiffEqFlux.basic_tgrad, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), Float32[2.0, 0.0], (0.0f0, 10.0f0), (layer_1 = NamedTuple(), layer_2 = (weight = Float32[-1.6285712 -0.7106906; 1.8252167 -0.15980828; … ; -0.46829274 0.3371172; -1.2275605 -1.2200289], bias = Float32[0.10678334, 0.04854786, -0.3036146, -0.67875606, 0.19386894, 0.5310734, -0.07256524, -0.18512556, -0.071528256, 0.37226313  …  -0.17194617, 0.21797575, -0.29705614, 0.19832478, 0.4268401, 0.43911377, 0.46748847, 0.19240205, 0.38869342, 0.29241425]), layer_3 = (weight = Float32[0.0004401929 0.031856854 … -0.056129575 0.2373046; -0.07581123 -0.13199675 … 0.20389898 -0.18144646], bias = Float32[-0.101620436, -0.1026898])), Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}(), SciMLBase.StandardODEProblem()), OrdinaryDiffEqTsit5.Tsit5{typeof(OrdinaryDiffEqCore.trivial_limiter!), typeof(OrdinaryDiffEqCore.trivial_limiter!), Static.False}(OrdinaryDiffEqCore.trivial_limiter!, OrdinaryDiffEqCore.trivial_limiter!, static(false)), OrdinaryDiffEqCore.InterpolationData{SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#dudt#17"{StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}, Vector{Float32}, Vector{Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}}, Nothing, OrdinaryDiffEqTsit5.Tsit5ConstantCache, Nothing}(SciMLBase.ODEFunction{false, SciMLBase.FullSpecialize, DiffEqFlux.var"#dudt#17"{StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}, LinearAlgebra.UniformScaling{Bool}, Nothing, typeof(DiffEqFlux.basic_tgrad), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}(DiffEqFlux.var"#dudt#17"{StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}}(StatefulLuxLayer{Static.True, Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}, Nothing, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}, layer_3::@NamedTuple{}}}(Chain{@NamedTuple{layer_1::WrappedFunction{Main.var"#1#2"}, layer_2::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = WrappedFunction(#1), layer_2 = Dense(2 => 50, tanh), layer_3 = Dense(50 => 2)), nothing), nothing, (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple()), nothing, static(true))), LinearAlgebra.UniformScaling{Bool}(true), nothing, DiffEqFlux.basic_tgrad, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, nothing, SciMLBase.DEFAULT_OBSERVED, nothing, nothing, nothing, nothing), CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}[Float32[2.0, 0.0], Float32[2.0003984, 0.084881306], Float32[2.000786, 0.16975331], Float32[2.0011172, 0.2546237], Float32[2.0013137, 0.33950973], Float32[2.0012627, 0.42445177], Float32[2.0008066, 0.5095357], Float32[1.9997556, 0.5948871], Float32[1.9978168, 0.6808015], Float32[1.9945722, 0.76778114]  …  Float32[-0.4347611, 0.30176672], Float32[-0.44770268, 0.3272455], Float32[-0.46191934, 0.35120934], Float32[-0.47744444, 0.37346473], Float32[-0.49423942, 0.39385107], Float32[-0.51219267, 0.41223925], Float32[-0.5311202, 0.42853218], Float32[-0.5507664, 0.44266585], Float32[-0.570801, 0.4546066], Float32[-0.5908468, 0.46434626]], Float32[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9  …  9.1, 9.2, 9.3, 9.4, 9.5, 9.6, 9.7, 9.8, 9.9, 10.0], Vector{CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}}[[Float32[2.0, 0.0]]], nothing, false, OrdinaryDiffEqTsit5.Tsit5ConstantCache(), nothing, false), false, 0, SciMLBase.DEStats(177, 0, 0, 0, 0, 0, 0, 0, 0, 0, 29, 0, 0.0), nothing, SciMLBase.ReturnCode.Success, nothing, nothing, nothing), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple()))

Neural ODE Example

Here is the full neural ODE example. Note that we use the gpu_device function so that the same code works on CPUs and GPUs, dependent on using LuxCUDA.

using Lux, Optimization, OptimizationOptimisers, Zygote, OrdinaryDiffEq, Plots, LuxCUDA,
      SciMLSensitivity, Random, ComponentArrays
import DiffEqFlux: NeuralODE

CUDA.allowscalar(false) # Makes sure no slow operations are occurring

#rng for Lux.setup
rng = Xoshiro(0)
# Generate Data
u0 = Float32[2.0; 0.0]
datasize = 30
tspan = (0.0f0, 1.5f0)
tsteps = range(tspan[1], tspan[2]; length = datasize)
function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u .^ 3)'true_A)'
end
prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
# Make the data into a GPU-based array if the user has a GPU
ode_data = solve(prob_trueode, Tsit5(); saveat = tsteps)
ode_data = Array(ode_data) |> gdev

dudt2 = Chain(x -> x .^ 3, Dense(2, 50, tanh), Dense(50, 2))
u0 = Float32[2.0; 0.0] |> gdev
p, st = Lux.setup(rng, dudt2)
p = p |> ComponentArray |> gdev
st = st |> gdev

prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(); saveat = tsteps)

predict_neuralode(p) = reduce(hcat, first(prob_neuralode(u0, p, st)).u)
function loss_neuralode(p)
    pred = predict_neuralode(p)
    loss = sum(abs2, ode_data .- pred)
    return loss
end
# Callback function to observe training
list_plots = []
iter = 0
callback = function (state, l; doplot = false)
    p = state.u
    global list_plots, iter
    pred = predict_neuralode(p)
    if iter == 0
        list_plots = []
    end
    iter += 1
    display(l)
    # plot current prediction against data
    plt = scatter(tsteps, Array(ode_data[1, :]); label = "data")
    scatter!(plt, tsteps, Array(pred[1, :]); label = "prediction")
    push!(list_plots, plt)
    if doplot
        display(plot(plt))
    end
    return false
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss_neuralode(x), adtype)
optprob = Optimization.OptimizationProblem(optf, p)
result_neuralode = Optimization.solve(
    optprob, OptimizationOptimisers.Adam(0.05); callback, maxiters = 300)

retcode: Default
u: ComponentArrays.ComponentVector{Float32, CUDA.CuArray{Float32, 1, CUDA.DeviceMemory}, Tuple{ComponentArrays.Axis{(layer_1 = ViewAxis(1:0, Shaped1DAxis((0,))), layer_2 = ViewAxis(1:150, Axis(weight = ViewAxis(1:100, ShapedAxis((50, 2))), bias = ViewAxis(101:150, Shaped1DAxis((50,))))), layer_3 = ViewAxis(151:252, Axis(weight = ViewAxis(1:100, ShapedAxis((2, 50))), bias = ViewAxis(101:102, Shaped1DAxis((2,))))))}}}(layer_1 = Float32[], layer_2 = (weight = Float32[-2.0023954 1.5009212; -0.47825912 -0.2519925; … ; -0.20992915 -0.0009007865; 0.5461162 -0.85644305], bias = Float32[-0.6268662, -0.7830921, -0.540858, 1.4213336, -0.40148732, 0.8359537, 0.8668236, -0.015425231, 0.95446044, 0.7596305  …  0.23153207, -0.9970776, -1.4622848, 0.5644652, -0.82362366, 0.71132743, -0.08877417, 1.0117041, 0.25556466, 0.3966038]), layer_3 = (weight = Float32[-0.2556873 0.7317984 … -0.05477792 0.6012788; -0.043002695 -0.06674853 … -0.13146661 -0.028954566], bias = Float32[-0.76554847, 0.18391721]))