python求高阶导数_Pythorch中的高阶梯度

我在pytorch中实现了以下Jacobian函数。除非我犯了错误，否则它计算任何张量w.r.t.的雅可比。任何维输入：import torch

import torch.autograd as ag

def nd_range(stop, dims = None):

if dims == None:

dims = len(stop)

if not dims:

yield ()

return

for outer in nd_range(stop, dims - 1):

for inner in range(stop[dims - 1]):

yield outer + (inner,)

def full_jacobian(f, wrt):

f_shape = list(f.size())

wrt_shape = list(wrt.size())

fs = []

f_range = nd_range(f_shape)

wrt_range = nd_range(wrt_shape)

for f_ind in f_range:

grad = ag.grad(f[tuple(f_ind)], wrt, retain_graph=True, create_graph=True)[0]

for i in range(len(f_shape)):

grad = grad.unsqueeze(0)

fs.append(grad)

fj = torch.cat(fs, dim=0)

fj = fj.view(f_shape + wrt_shap