Как эффективно использовать автоград PyTorch с тензорами?

В моем предыдущем вопросе я узнал, как использовать автоградиент 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) 

        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()
r = torch.tensor([1.0], requires_grad=True)
print('r', r)
y = nx(r)
print('y', y)
print('')
for i in range(y.shape[0]):
    # prints the vector (dy_i/dr_0, dy_i/dr_1, ... dy_i/dr_n)
    print(grad(y[i], r, retain_graph=True))

>>>
r tensor([1.], requires_grad=True)
y tensor([ 0.1698, -0.1871, -0.1313, -0.2747], grad_fn=<AddBackward0>)

(tensor([-0.0124]),)
(tensor([-0.0952]),)
(tensor([-0.0433]),)
(tensor([-0.0099]),)

Проблема, с которой я сейчас сталкиваюсь, заключается в том, что мне нужно различать очень большой тензор, и повторение его, как я делаю сейчас (for i in range(y.shape[0])), занимает вечность. Причина, по которой я повторяю, заключается в том, что, исходя из понимания, grad знает только, как распространять градиенты из скалярного тензора, а y - нет. Поэтому мне нужно вычислить градиенты по каждой координате y.
Я знаю, что TensorFlow способен различать тензоры от здесь:

tf.gradients(
    ys, xs, grad_ys=None, name='gradients', gate_gradients=False,
    aggregation_method=None, stop_gradients=None,
    unconnected_gradients=tf.UnconnectedGradients.NONE
)
"ys and xs are each a Tensor or a list of tensors. grad_ys is a list of Tensor, holding the gradients received by the ys. The list must be the same length as ys.

gradients() adds ops to the graph to output the derivatives of ys with respect to xs. It returns a list of Tensor of length len(xs) where each tensor is the sum(dy/dx) for y in ys and for x in xs."

И надеялся, что в PyTorch есть более эффективный способ различать тензоры.

Например:

a = range(100)
b = range(100)
c = range(100)
d = range(100)
my_tensor = torch.tensor([a,b,c,d])

t = range(100)

#derivative = grad(my_tensor, t) --> not working

#Instead what I'm currently doing:
for i in range(len(t)):
    a_grad = grad(a[i],t[i], retain_graph=True)
    b_grad = grad(b[i],t[i], retain_graph=True)
    #etc.

Мне сказали, что это могло бы сработать, если бы я мог запустить autograd на прямом проходе, а не на обратном, но из здесь похоже, что в настоящее время PyTorch не поддерживает эту функцию.

Обновление 1:
@jodag упомянул, что то, что я ищу, может быть просто диагональю якобиана. Я перехожу по ссылке, которую он прикрепил, и пробую более быстрый метод. Хотя это, похоже, не работает и выдает ошибку: RuntimeError: grad can be implicitly created only for scalar outputs. Код:

nx = net_x()
x = torch.rand(10, requires_grad=True)
x = torch.reshape(x, (10,1))
x = x.unsqueeze(1).repeat(1, 4, 1)
y = nx(x)
dx = torch.diagonal(torch.autograd.grad(torch.diagonal(y, 0, -2, -1), x), 0, -2, -1)