In our case, we have to minimize the cost function \( c(x, P) \) with respect to the two sets of weights and bisases, that is for the hidden layer \( P_{\mathrm{hidden} } \) and for the ouput layer \( P_{\mathrm{output} } \) .
This means that \( P_{\mathrm{hidden} } \) and \( P_{\mathrm{output} } \) is updated by $$ \begin{align} P_{\mathrm{hidden},\mathrm{new}} &= P_{\mathrm{hidden}} - \lambda \nabla_{P_{\mathrm{hidden}}} c(x, P) \tag{20}\\ P_{\mathrm{output},\mathrm{new}} &= P_{\mathrm{output}} - \lambda \nabla_{P_{\mathrm{output}}} c(x, P) \tag{21} \end{align} $$
This might look like a cumberstone to set up the correct expression for finding the gradients. Luckily, Autograd comes to the rescue.
def solve_ode_neural_network(x, num_neurons_hidden, num_iter, lmb):
## Set up initial weigths and biases
# For the hidden layer
p0 = npr.randn(num_neurons_hidden, 2 )
# For the output layer
p1 = npr.randn(1, num_neurons_hidden + 1 ) # +1 since bias is included
P = [p0, p1]
print('Initial cost: %g'%cost_function(P, x))
## Start finding the optimal weigths using gradient descent
# Find the Python function that represents the gradient of the cost function
# w.r.t the 0-th input argument -- that is the weights and biases in the hidden and output layer
cost_function_grad = grad(cost_function,0)
# Let the update be done num_iter times
for i in range(num_iter):
# Evaluate the gradient at the current weights and biases in P.
# The cost_grad consist now of two arrays;
# one for the gradient w.r.t P_hidden and
# one for the gradient w.r.t P_output
cost_grad = cost_function_grad(P, x)
P[0] = P[0] - lmb * cost_grad[0]
P[1] = P[1] - lmb * cost_grad[1]
print('Final cost: %g'%cost_function(P, x))
return P