Hi everyone. I’m trying to run a FNN in which the derivative of the output is used in the loss function, as well as functions that have the output of the whole NN as their arguments.

Basically, I am struggling on how to make everything understandable for Flux. I have spent quite a time reading Flux tutorials, searching for similar questions and even trying to make parallels with similar problems being solved in Python, however my code doesn’t run.

I will replicate here and try to make sense of what I’m trying to do, so perhaps some of you can kindly give me some light.

```
using Flux
using ReverseDiff
##Train Data
t=vcat(0:0.1:4)
out=zeros(size(t)[1])
##Problem parameters
α=1; C=0.5; β=2; P=5; π0=1
##Building the model
len_hidden=5 #length of hidden layer
x=Chain(Dense(1,len_hidden,σ),Dense(len_hidden,3,relu))
dxdt(x)=x->ReverseDiff.gradient(x)
p=Chain(Dense(1,len_hidden,σ),Dense(len_hidden,3,relu))
dpdt(p)=p->ReverseDiff.gradient(p)
u=Chain(Dense(1,len_hidden,σ),Dense(len_hidden,3,relu))
#H
H(x,p,u)=α*u*x-C*u^2+p*β*x*(1-x)*(P*u-π0)
#Partials
dHdx(H,x)=gradient(H,x,p,u)[1]
dHdp(H,p)=gradient(H,x,p,u)[2]
dHdu(H,u)=gradient(H,x,p,u)[3]
#Full Model
#The last two terms are just impositions on initial condition for x and final condition for p
model(x,p,u)=(dxdt+dHdp)^2+(dpdt-dHdx)^2+(dHdu)^2+(x(0)-0.10)^2+(p(4)-1)^2
#Training
opt=Descent()
parameters=params(x,p,u)
data=[(t,out)]
loss(model,out)=Flux.Losses.mse(model(x,p,u),out)
Flux.train!(loss, parameters, data, opt)
```

My single hidden layer NN should have time as input and spit out x,p and u, all functions of time.

The loss function needs the time derivatives of each output (therefore the derivative with respect to input - which I am simply trying to tackle by the ReverseDiff.gradient). It also needs the partials of the function H, which is itself a function of the outputs as well. H=H(x,p,u).

I have been changing, tuning and trying to understand how to correctly implement this for a a while.

General observations, tips, comments are appreciated. Even if you think I should consider another package or approach for my problem, please feel free to let me know.

Thank you for you collaboration.

Gabriel