Batch evaluation for NeuralODE?

Specific Domains Machine Learning
1

Hi, I am new to Julia and am trying to use NeuralODE as follows:

d_in, dim_hidden = 25, 16
tspan = (0., 1.)
t = range(tspan[1], tspan[2], length=100)
dxdt = Chain(Dense(d_in, dim_hidden, tanh), Dense(dim_hidden, dim_hidden, tanh), Dense(dim_hidden, d_in,tanh))
n_ode = NeuralODE(dxdt, tspan, Tsit5(), saveat=t, reltol=1e-7, abstol=1e-9)

I found that when I pass a single array of shape (25, ), the model works OK:

tmp_in = rand(Float64, (25))
result = n_ode(tmp_in)

But if I pass a batch of arrays which has shape (115, 25), then it would not work:

tmp_in = rand(Float64, (115, 25))
result = n_ode(tmp_in)

and gives following error:

DimensionMismatch("A has dimensions (16,25) but B has dimensions (115,25)")

I guess it is more like a Dense layer issue, but am not exactly sure. Does anyone have idea how I should fix it? Thanks!

2

It’s saying you’re multiplying a 16x25 matrix to a 115x25 matrix. Batches are columns, did you mean tmp_in = rand(Float64, (25, 115))?