In Pytorch, I have a multi-dimension solution stored on each column of a tensor
u, and for the training of the PINN I’m developing, I use this function to compute the derivative column-wise.
def dt(self, u, t): # return batch_jacobian(u, t) N = u.shape u_t =  # initializing the time derivatives # for each columns (i), we compute the time derivative ui_t for i in range(N): ui_t = torch.autograd.grad(u[:, i], t, grad_outputs=torch.ones_like(u[:, i]), retain_graph=True, create_graph=True, allow_unused=True) u_t.append(ui_t) # time derivatives are stored in u_t u_t = torch.cat(tuple(u_t), dim=1) # we concateneate all the derivatives return u_t
but it involves a
for loop that I would like to remove.
In my research about it, I found this topic with the same issue, but I wasn’t able to fix my code with the proposed solutions.
Here is the definition of the functional that is used for the training of the neural network :
def functional_f(self, t, theta): self.A = torch.tensor(theta[:,1:], requires_grad=False).float().to(self.device) self.mu = torch.tensor(theta[:,0], requires_grad=False).float().to(self.device) u = self.dnn(t) u_t = self.dt(u, t) ode = u_t - self.renorm*(self.mu + (self.A @ torch.exp(u).T).T) return ode
the ODE system to be solved is ∂t ui = µ + A⋅exp(ui), where µ ∈ ℝ^N and A ∈ ℝ^(N,N) are stored in the matrix
theta, are fixed parameters involved in the ODE.
How could I implement the function
dt so that there is no more loop in it ?
I also tried the decorator
jit but without success either…
(nb : this is a duplicate of my stack overflow post that did not get an answer yet…)