PyTorch计算4D矩阵Jacobian结果异常的原因排查

1. Jacobian输出的维度结构与预期不匹配 #

torch.autograd.functional.jacobian的输出规则是：将输出张量展平为一维数组，再计算每个展平元素对输入张量的偏导。对于你的4×4矩阵输出，返回的Jacobian形状为(16, 4)（16个矩阵元素 × 4个输入变量），而非直观的(4,4,4)（矩阵行×矩阵列×输入维度）。

如果直接将(16,4)的结果按默认方式reshape为(4,4,4)，或者错误地解读维度顺序，就会出现非零项位置错位、不该有值的区域出现非零的情况。

正确的维度转换方式是将结果转置后reshape，得到每个输入变量对应的偏导矩阵：

jacobian = torch.autograd.functional.jacobian(metric, coord)
# 转换为 (输入变量数, 矩阵行, 矩阵列)，即每个切片对应一个输入变量的偏导矩阵
jacobian_per_input = jacobian.T.reshape(4, 4, 4)

此时：

• jacobian_per_input[0,:,:]对应对t的偏导，全为0，符合解析解
• jacobian_per_input[1,:,:]对应对r的偏导，仅(2,2)和(3,3)位置有非零值（2r和2r sin²θ）
• jacobian_per_input[2,:,:]对应对theta的偏导，仅(3,3)位置有非零值（2r² sinθ cosθ）
• jacobian_per_input[3,:,:]对应对phi的偏导，全为0，符合解析解

2. 次要问题：数值类型混用（不影响位置，但可能引发精度问题） #

你的代码中使用了np.pi来定义theta输入，但输入张量是torch.tensor，建议替换为torch.pi，避免numpy与PyTorch数值类型隐式转换带来的潜在精度问题：

test_minkowski_jacob(minkowski_metric_spherical, 0.0, 1.0, torch.pi / 4, 0.0)

验证解析解的正确维度对应 #

回顾Minkowski球坐标度量的解析偏导：

• $g_{00}=-1$、$g_{11}=1$：对所有输入变量偏导为0
• $g_{22}=r²$：仅对r的偏导为2r
• $g_{33}=r² sin²θ$：对r的偏导为2r sin²θ，对theta的偏导为2r² sinθ cosθ，其余偏导为0

经过正确维度转换后的Jacobian结果完全符合上述解析解。

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