通过计算雅可比行列式有效地使用PyTorch的autograd与张量

在我之前的问题中，我发现了如何将 PyTorch 的 autograd 与张量一起使用：

import torch
from torch.autograd import grad
import torch.nn as nn
import torch.optim as optim

class net_x(nn.Module): 
        def __init__(self):
            super(net_x, self).__init__()
            self.fc1=nn.Linear(1, 20) 
            self.fc2=nn.Linear(20, 20)
            self.out=nn.Linear(20, 4) #a,b,c,d

        def forward(self, x):
            x=torch.tanh(self.fc1(x))
            x=torch.tanh(self.fc2(x))
            x=self.out(x)
            return x

nx = net_x()

#input
t = torch.tensor([1.0, 2.0, 3.2], requires_grad = True) #input vector
t = torch.reshape(t, (3,1)) #reshape for batch

#method 
dx = torch.autograd.functional.jacobian(lambda t_: nx(t_), t)
dx = torch.diagonal(torch.diagonal(dx, 0, -1), 0)[0] #first vector
#dx = torch.diagonal(torch.diagonal(dx, 1, -1), 0)[0] #2nd vector
#dx = torch.diagonal(torch.diagonal(dx, 2, -1), 0)[0] #3rd vector
#dx = torch.diagonal(torch.diagonal(dx, 3, -1), 0)[0] #4th vector
dx 
>>> 
tensor([-0.0142, -0.0517, -0.0634])

问题是 grad只知道如何从标量张量传播梯度（我的网络输出不是），这就是我必须计算雅可比行列式的原因。

但是，这不是很有效并且有点慢，因为我的矩阵很大并且计算整个雅可比矩阵需要一段时间（而且我也没有使用整个雅可比矩阵）。

有没有办法只计算雅可比矩阵的对角线（在这个例子中得到 4 个向量）？

似乎有一个开放的功能请求但似乎没有得到太多关注。

更新 1：
我尝试了@iacob 所说的设置torch.autograd.functional.jacobian(vectorize=True)。
然而，这似乎更慢。为了测试这一点，我将网络输出从 更改4为400，并将输入t更改为：

val = 100
t = torch.rand(val, requires_grad = True) #input vector
t = torch.reshape(t, (val,1)) #reshape for batch

没有vectorized = True：

Wall time: 10.4 s

和：

Wall time: 14.6 s