Speed up Jacobian matrix calculation in Pytorch

username: S_M
cooked: <p>Hi all<br>
I use ‘torch.autograd.functional.jacobian(f,x)’ to calculate the partial derivatives of f with respect to x but when the dimension of f is incremented the time of calculation increases.<br>
does anyone know any solution to speed up the jacobian calculation in PyTorch?</p>
<p>Thanks in advance.</p>

username: AlphaBetaGamma96
cooked: <p>You might want to have a look at <a href="https://github.com/pytorch/functorch/blob/main/README.md" rel="noopener nofollow ugc">FuncTorch</a></p>

username: S_M
cooked: <aside class="quote no-group" data-username="AlphaBetaGamma96" data-post="2" data-topic="143475">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/3da27b/48.png" class="avatar"> AlphaBetaGamma96:</div>
<blockquote>
<p>You might want to hav</p>
</blockquote>
</aside>
<p>It has an issue…<br>
Successfully built functorch<br>
Installing collected packages: functorch<br>
Successfully installed functorch-0.2.0a0+a2d2f99<br>
but when I run my code<br>
from functorch import vmap, grad<br>
ModuleNotFoundError: No module named ‘functorch’</p>

username: AlphaBetaGamma96
cooked: <p>Did you install it within the current environment? For example, running pytorch in pip but installing functorch with conda?</p>

username: S_M
cooked: <p>I installed the package in my current environment using pip command as had been described in the installation guide. But, the name of my current environment is test.</p>

username: AlphaBetaGamma96
cooked: <p>If you’re using Linux, you can check what python3 executable you’re using with <code>which python</code>, if you installed <code>FuncTorch</code> you most likely made a new environment (as is stated in the install guide) and installed it in that. Also check you didn’t install <code>pytorch</code> with conda when you installed! Then if that’s all ok, manually check your environemnt and see where <code>pytorch</code> is installed and where <code>functorch</code> is installed?</p>

username: S_M
cooked: <p>It doesn’t work on my local machine. While installing the functorch package some errors occur.</p>

username: AlphaBetaGamma96
cooked: <p>Did you follow the instruction of the <a href="https://colab.research.google.com/drive/1GNfb01W_xf8JRu78ZKoNnLqiwcrJrbYG#scrollTo=HJ1srOGeNCGA" rel="noopener nofollow ugc">google colab</a>? When you install functorch in a new environment you’ll need to ensure you’ve PyTorch (among other dependencies) reinstalled for the new environment too.</p>

username: S_M
cooked: <p>Yes, I followed them without any problem but when my google drive was mounted to load my code which used the’ functorch’ package the error of there is no package by the name of the functorch was shown.</p>

username: AlphaBetaGamma96
cooked: <p>Try installing after mounting your google drive? It sounds like you’re using different environments.</p>

username: S_M
cooked: <p>As I checked <a href="https://colab.research.google.com/drive/1CrLkqIrydBYP_svnF89UUO-aQEqNPE8x?usp=sharing" rel="noopener nofollow ugc">google colab guid</a> I don’t need to create an environment. Am I right?</p>

username: S_M
cooked: <p>I mounted my google drive at the first then the functorch package was installed. Finally, I tried to run my code by the colab unfortunately an error has occurred.</p>
<p>tcmalloc: large alloc 140349731364864 bytes == (nil) @  0x7fa5b927e887 0x7fa5b7b74c29 0x7fa5b7b75afb 0x7fa5b7b7620c 0x7fa4f64e6418 0x7fa4f64ea2d8 0x7fa4aee1bd1d 0x7fa4aee151b1 0x7fa5b94c98d3 0x7fa5b94ce39f 0x7fa5b81ac1ef 0x7fa5b94cd96a 0x7fa5b8e26f96 0x7fa5b81ac1ef 0x7fa5b81ac27f 0x7fa5b8e27745 0x7fa5b8e27051 0x7fa5ab7a3a16 0x55d25a847317 0x55d25a847120 0x55d25a8baf33 0x55d25a8b666e 0x55d25a849151 0x55d25a88a7b9 0x55d25a8876d4 0x55d25a847b32 0x55d25a8bba6d 0x55d25a8489da 0x55d25a8bb2c0 0x55d25a8b602f 0x55d25a8b5d43<br>
terminate called after throwing an instance of ‘std::bad_alloc’<br>
what():  std::bad_alloc</p>
<p>Do you have any advice?</p>

username: AlphaBetaGamma96
cooked: <p>Looks like you’re tring to allocate a lot of RAM. It’s most likely a common error on google-drive and I’d check on stackexchange to see if it’s been noted there!</p>

username: S_M
cooked: <p>I’d check it by using batch 8! it got the error again!!!</p>