Ok, thank you, I see. So in this case, I would have a couple of follow-up questions.
If I understand you correctly, I would have to create dy/dx[k,i,j] step by step:


This would involve taking the gradient of y[k] with respect to every element of x[i,j].
If I define x as in the above code, x[i,j].grad is None. If I manually create x11,x12 … and then construct the x from it using stack, then I can get the gradient, but it is very gawky. Is there a way to define x as a matrix but then compute the derivative of each individual matrix item?


Would I then call y[k].backward() every time before looping through all (i,j) to get the gradients?


Updated code with the questions in the comments:
import torch

n, p, q = 9, 2, 3

# this is the gawky part
x11 = torch.tensor(1.1815, requires_grad=True)
x12 = torch.tensor(-0.2189, requires_grad=True)
x13 = torch.tensor(-0.9923, requires_grad=True)
x21 = torch.tensor(0.2705, requires_grad=True)
x22 = torch.tensor(0.9501, requires_grad=True)
x23 = torch.tensor(-0.4003, requires_grad=True)
X_lst = [[x11, x12, x13], [x21,x22,x23]]

x1 = x11
x1 = torch.hstack((x1, x12))
x1 = torch.hstack((x1, x13))
x2 = x21
x2 = torch.hstack((x2, x22))
x2 = torch.hstack((x2, x23))
x  = torch.vstack((x1,x2))
# just to get the x

A = torch.randn(q, p)
z = torch.randn(q, 1)

y = None
for k in range(int(n/q)):
    if(y is None):
        f = A@x
        y = f @ z
    else:
        f = A@x @ f
        y = torch.vstack( ( y, f @ z ) )

# need to do this? here?
# for k in range(n):
#     y[k].backward(retain_graph=True)

dY_dX = torch.zeros((n,p,q))
for i in range(p):
    for j in range(q):
        for k in range(n):
            # or do it here?
            y[k].backward(retain_graph=True)
            xx = X_lst[i][j]
            xx.retain_grad()
            dyk_dxij = xx.grad.clone()
            xx.grad.zero_()

            # because x[i,j].grad gives me None
            # dyk_dxij = x[i,j].grad

            dY_dX[k, i, j] = dyk_dxij


print('dydx with autograd: \n', dY_dX)

Is this the way to go?