Error with training decoupled NODEs

I’m attempting to train two decoupled neural ODEs simultaneously. Although decoupled, these ODEs interact through a coupled loss function (not shown). I have constructed this in two ways: 1) by writing two separate ODEProblem() definitions and evaluating the dependent loss function and 2) as included in this post; that is, by writing a single ODEProblem() definition as a decoupled pair of ODEs. Both methods fail with strange errors.

A minimum working example (I’ve removed all of the training data and custom loss function) is as follows (I’m running Julia 1.6):

using DiffEqFlux, DifferentialEquations, Flux

tSpan = (0.0,3.0)
tSteps = range(tSpan[1], tSpan[2], length=4)

L = FastChain(FastDense(1,3,sigmoid),FastDense(3,1))
p = [initial_params(L);-0.1] # appended w/ initial param for linear ODE to be trained

function dudt(u,p,t)
	z = L(u[1],p[1:end-1])
	return [p[end]*u[1],z[1]]
end

evolveODE = ODEProblem(dudt, [1.0,1.0], tSpan, p)

function lossFun(p)

	u = Array(solve(evolveODE, Tsit5(), saveat = tSteps))
	loss = sum(abs2, u[1,:])

	return loss, u[1,:]
end


callback = function (p, l, pred)
  display(l)
  return false
end

result_neuralode = DiffEqFlux.sciml_train(lossFun,
                                           p,
                                           ADAM(0.05),
                                           cb = callback,
                                           maxiters = 5)

As written, the ODEProblem() can be solved straight away in the REPL, However, in this simple script, it fails with the error:

MethodError: Cannot convert an object of type Bool to an object of type Matrix{Float64}

This confuses me - what Boolean object exists in this example? As mentioned, I tried this two different ways - the other method also gives me a different error, but this example is simpler to post and work through.

Any thoughts?

The neural network needs an array input. Working code:

using DiffEqFlux, DifferentialEquations, Flux

tSpan = (0.0,3.0)
tSteps = range(tSpan[1], tSpan[2], length=4)

L = FastChain(FastDense(1,3,sigmoid),FastDense(3,1))
p = [initial_params(L);-0.1] # appended w/ initial param for linear ODE to be trained

function dudt(u,p,t)
	z = L([u[1]],p[1:end-1])
	return [p[end]*u[1],z[1]]
end

evolveODE = ODEProblem(dudt, [1.0,1.0], tSpan, p)

function lossFun(p)

	u = Array(solve(evolveODE, Tsit5(), p=p, saveat = tSteps))
	loss = sum(abs2, u[1,:])

	return loss, u[1,:]
end


callback = function (p, l, pred)
  display(l)
  return false
end

result_neuralode = DiffEqFlux.sciml_train(lossFun,
                                           p,
                                           ADAM(0.05),
                                           cb = callback,
                                           maxiters = 500)