NumPy中矩阵对自身雅可比矩阵（2D扩展4D）的原生实现方法查询

开发者阿泽

2026-10-17

解决方案

方法1：reshape+对角赋值（最优性能）

原理是将目标4D数组展开为(k*d, k*d)的2D数组，其对角线元素恰好对应你需要赋值的(i,j,i,j)位置，直接用NumPy原生的np.fill_diagonal实现赋值，完全避免Python级循环：

import numpy as np

def mat_to_4d_diag(a):
    k, d = a.shape
    r = np.zeros((k, d, k, d), dtype=a.dtype)
    # 展开为2D后填充对角线
    np.fill_diagonal(r.reshape(k*d, k*d), a.ravel())
    return r

该实现仅用到NumPy原生接口，完全符合你仅使用NumPy原生函数开发的要求，所有运算都在NumPy底层C层面执行，相比循环实现速度提升可达数百到数千倍，矩阵越大性能提升越明显。

方法2：广播乘法实现

如果你更倾向于用张量运算的逻辑实现，可以利用单位矩阵的广播相乘得到同样结果：

def mat_to_4d_diag_v2(a):
    k, d = a.shape
    # 构造两个维度对应的单位矩阵，广播相乘后和a的扩展维度相乘
    return a.reshape(k,d,1,1) * np.eye(k).reshape(k,1,k,1) * np.eye(d).reshape(1,d,1,d)

该方法性能略逊于方法1，但逻辑更贴合张量变换的直观理解，适合需要结合其他张量运算的场景。