Speed up Jacobian in PyTorch

[caowen1510957]

Hello, I saw that here is a place to leave suggestions and use cases.

I am working on solving PDEs (numerical method) using PyTorch, I cannot avoid Jacobian (Newton iteration, used for solving nonlinear equations), namely, the Jacobian matrix of residual vector. But it is too slow. Here is the runtime (get_residual is the function I wrote to construct the equation I need to solve, and Jacobian is the built-in function, imported from PyTorch):

After I used Functorch, the performance is way better

I don't know if there is more room to speed up, attach is my code as the use case, when you increase Nx, Ny, and Nz, the runtime is crazy. main.txt