Get the derivatives after training

[phdstudentjw]

Suppose I already finished training a neural network to approximate a PDE called u. Usually, I would use model.predict(X) to obtain values of u over domain X. However, how can I get the derivatives of u with respect to x? I know how to use gradients.jacobian to define a loss function before training, but I don't know how I can get the jacobian after training. Thank you. I looked into state_dict() but I don't think it has what I want.

Could you please help? Thank you.

[riccardotomada]

Suppose your solution is called 'y' and its derivative wrt the first coordinate is 'u'. Then you can define a function that computes u. Finally you only have to assign this function to the 'operator' parameter inside model.predict()

def get_u(x, y):
   return dde.grad.jacobian(y, x, i = 0, j = 1) 

u = model.predict(X, operator = get_u)