使用PyTorch/JAX验证矩阵函数导数张量的实现问题排查

针对函数 $f(x) = \frac{Axx^\top A}{x^\top A^2 x}$（$A$ 为对称半正定方阵，$x$ 为向量）的导数验证，以下是可运行的PyTorch和JAX实现，同时说明常见错误原因：

常见错误排查 #

• 直接用torch.grad/jax.grad求导：这两个API默认仅支持标量输出的函数求导，而$f(x)$是矩阵值函数，必须通过Jacobian相关方法处理。
• 向量维度错误：未将$x$转为列向量会导致矩阵乘法维度不匹配，引发报错。
• 分母维度未处理：$x^\top A^2 x$是标量，需挤压多余维度避免广播错误。

PyTorch 实现代码 #

import torch

def f(x, A):
    x = x.unsqueeze(1)  # 将向量转为列向量 (n, 1)
    numerator = A @ x @ x.T @ A
    denominator = x.T @ A @ A @ x  # 计算标量分母 x^T A² x
    return numerator / denominator.squeeze()  # 挤压分母的多余维度

def compute_derivative(x, A):
    x = x.detach().requires_grad_(True)
    output = f(x, A)
    grad_tensor = []
    # 遍历输出矩阵的每个元素，计算其对x的梯度，组合为三阶张量
    for i in range(output.shape[0]):
        row_grads = []
        for j in range(output.shape[1]):
            grad = torch.autograd.grad(output[i,j], x, retain_graph=True)[0]
            row_grads.append(grad)
        grad_tensor.append(torch.stack(row_grads, dim=1))
    return torch.stack(grad_tensor, dim=0)

# 测试用例
n = 3
# 构造对称半正定矩阵A
A = torch.randn(n, n)
A = A @ A.T
x = torch.randn(n)

# 计算导数张量（形状为 (n, n, n)）
deriv = compute_derivative(x, A)
print("导数张量形状:", deriv.shape)

# 有限差分验证
h = torch.randn(n) * 1e-6
f_x = f(x, A)
f_xh = f(x + h, A)
finite_diff = (f_xh - f_x) / h.norm()
# 导数作用于向量h，得到(n,n)矩阵
deriv_h = torch.einsum('ijk,k->ij', deriv, h)
print("自动微分与有限差分的误差:", torch.norm(deriv_h - finite_diff))

JAX 实现代码 #

import jax
import jax.numpy as jnp

def f(x, A):
    x = x[:, jnp.newaxis]  # 转为列向量 (n, 1)
    numerator = A @ x @ x.T @ A
    denominator = x.T @ A @ A @ x
    return numerator / denominator.squeeze()

# 用jax.jacrev计算输出对x的Jacobian，直接得到三阶张量
compute_derivative = jax.jacrev(f, argnums=0)

# 测试用例
n = 3
key = jax.random.PRNGKey(42)
# 构造对称半正定矩阵A
A = jax.random.normal(key, (n, n))
A = A @ A.T
x = jax.random.normal(key, (n,))

# 计算导数张量
deriv = compute_derivative(x, A)
print("导数张量形状:", deriv.shape)

# 有限差分验证
h = jax.random.normal(key, (n,)) * 1e-6
f_x = f(x, A)
f_xh = f(x + h, A)
finite_diff = (f_xh - f_x) / jnp.linalg.norm(h)
deriv_h = jnp.einsum('ijk,k->ij', deriv, h)
print("自动微分与有限差分的误差:", jnp.linalg.norm(deriv_h - finite_diff))

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