Градиенты высшего порядка в pytorch

Я реализовал следующую функцию Якоби в pytorch. Если я не ошибся, он вычисляет якобиан любого тензора относительно любые размерные входы:

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_shape)
    return fj

Вдобавок к этому я попытался реализовать рекурсивную функцию для вычисления производных n-го порядка:

def nth_derivative(f, wrt, n):
    if n == 1:
        return full_jacobian(f, wrt)
    else:        
        deriv = nth_derivative(f, wrt, n-1)
        return full_jacobian(deriv, wrt)

Я провел простой тест:

op = torch.ger(s, s)
deep_deriv = nth_derivative(op, s, 5)

К сожалению, мне удалось получить гессиан ... но не производные более высокого порядка. Я знаю, что многие производные более высокого порядка должны быть равны 0, но я бы предпочел, чтобы pytorch мог это вычислить аналитически.

Одно из исправлений заключалось в изменении расчета градиента на:

try:
            grad = ag.grad(f[tuple(f_ind)], wrt, retain_graph=True, create_graph=True)[0]
        except:
            grad = torch.zeros_like(wrt)

Это общепринятый правильный способ справиться с этим? Или есть вариант получше? Или у меня изначально совершенно неправильная причина моей проблемы?