带 Torch 的矩阵偏导数

Question

我有一个函数 f(t,k)=tk+2tk^2。我想得到关于 t=[1,2] 和 k=[2,3] 的 k 的偏导数。我创建了一个张量矩阵，如下所示： 伽玛= [10,21 20,42]

当我尝试获得关于 k 的 gamma 的导数时，我期望矩阵 2*2 但 torch 返回一个向量。有什么问题吗？

import torch
import numpy as np

t=torch.tensor([1.0], requires_grad=True)
k=torch.tensor([1.0], requires_grad=True)
gamma=torch.zeros((2,2))

def f(a,b): #c is a coeff vector
    global t,k,gamma
    t=torch.tensor(a, dtype=float,requires_grad=True)
    k=torch.tensor(b, dtype=float,requires_grad=True)
    for i in range(2):
        for j in range(2):
            gamma[i,j]=torch.mul(t[i],k[j])+torch.mul(2,torch.mul(t[i],torch.pow(k[j],2)))

a=[1,2]
b=[2,3]
f(a,b)

def dy_dx(y, x):
    return torch.autograd.grad(y, x, grad_outputs=torch.ones_like(y),allow_unused=True, create_graph=True)
    
print(dy_dx(gamma,k))

我试图获得 f(t,k) 对 k 的导数，其中 t=[1,2] 和 k=[2,2]。我期望矩阵如下： [9,13 18,26]

但是 torch 返回一个向量，如下所示： [27., 39.]