PyTorch高阶梯度计算问题：递归求n阶导数遇阻

首先，你的问题根源在于当高阶导数为常数（或零）时，对应的张量元素不再具有可导性。比如你测试的torch.ger(s,s)是二次函数，三阶及以上导数都是零，此时f[tuple(f_ind)]是一个常数张量，PyTorch无法对其求导（因为常数没有梯度信息），从而抛出异常。

你的try-except方案的评价 #

用try-except捕获异常并返回零张量是一种可行的临时解决办法，但它存在两个明显的问题：

• 不够精准：会捕获所有类型的异常（比如索引错误、张量形状不匹配等），可能掩盖其他潜在的bug。
• 不够优雅：依赖异常处理来实现正常逻辑，虽然符合Python的"EAFP"风格，但在这里属于不得已的权宜之计，不如主动判断逻辑清晰。

更优的解决方案 #

1. 提前判断可导性，避免异常 #

在调用ag.grad之前，先检查当前张量元素是否可导。如果是常数（没有梯度函数），直接返回零张量，这样可以避免触发异常，逻辑也更清晰：

修改full_jacobian中的梯度计算部分：

import torch
import torch.autograd as ag
from itertools import product

def full_jacobian(f, wrt):
    f_shape = list(f.size())
    wrt_shape = list(wrt.size())
    fs = []
    # 用itertools.product替代递归的nd_range，更高效简洁
    f_range = product(*map(range, f_shape))
    
    for f_ind in f_range:
        f_element = f[tuple(f_ind)]
        # 判断当前元素是否可导：requires_grad为True且存在梯度函数
        if f_element.requires_grad and f_element.grad_fn is not None:
            grad = ag.grad(f_element, wrt, retain_graph=True, create_graph=True)[0]
        else:
            grad = torch.zeros_like(wrt)
        # 添加对应维度的unsqueeze
        for _ 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

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)

2. 优化索引生成逻辑 #

你原来的nd_range递归生成器可以用itertools.product替代，代码更简洁且运行效率更高——毕竟itertools的实现是C级别的，比手动递归生成器快不少。

3. 验证高阶导数结果 #

修改后，测试你的例子：

s = torch.tensor([1.0, 2.0], requires_grad=True)
op = torch.ger(s, s)
deep_deriv = nth_derivative(op, s, 5)
print(deep_deriv.shape)  # 输出 (2,2,2,2,2,2,2)，对应2x2的输出，5次对2维输入求导的形状
print(torch.allclose(deep_deriv, torch.zeros_like(deep_deriv)))  # 输出 True，验证全零结果

此时三阶及以上导数都会正确返回全零张量，不会抛出异常。

额外注意事项 #

• 内存与计算图：递归求高阶导数会构建非常深的计算图，retain_graph=True会保留所有中间计算图，对于大张量可能导致内存占用过高。如果只需要数值结果而非计算图，可以在最后一次求导时设置create_graph=False，或者在合适的时机调用torch.no_grad()。
• 官方API替代：PyTorch 1.10+提供了torch.autograd.functional.jacobian，官方实现更高效且经过优化，可以直接替代你手动实现的full_jacobian：

  import torch.autograd.functional as af
  def full_jacobian(f, wrt):
      # 注意：af.jacobian返回的形状是f_shape + wrt_shape，和你的实现完全一致
      return af.jacobian(f, wrt, create_graph=True)
  

  用官方函数替代后，递归逻辑可以保持不变，且能避免很多手动实现的bug。

内容的提问来源于stack exchange，提问作者user650261