基于NumPy实现兼容多维度张量的通用导数计算方法

完全可行。你当前分场景实现的本质是计算逐元素函数$f(x)=x^2$的雅可比张量，这类逐元素运算的导数天然具备对角稀疏性，不需要针对不同输入维度写分支逻辑，也不需要强制把输入规整为二维数组，完全可以写出适配0维标量、1维向量、2维矩阵乃至任意高阶张量的通用实现，且性能远高于多层嵌套循环的写法。

你三个独立函数的导数规律是完全统一的：

• 0维标量输入：导数为0维张量，取值为2*x，形状和输入一致
• 1维长度为n的向量输入：导数为n×n的对角矩阵，仅对角位置取值为对应位置的2*x[i]，其余位置为0
• 2维形状为(m,n)的矩阵输入：导数为形状(m,n,m,n)的四阶张量，仅当输出位置索引和输入位置索引完全相等时取值为2*x[i,j]，其余位置为0

推广到任意k维张量输入，导数张量的形状永远是x.shape + x.shape，且仅所有维度索引一一对应的对角位置有非零值2*x，其余位置均为0。

不需要写任何输入维度判断分支，也不需要多层循环，通过展平重构+对角填充的方式即可实现全维度适配：

import numpy as np

def foo(x):
    x = np.asarray(x)
    # 计算函数值，逐元素平方自动适配任意维度
    f = x ** 2
    x_size = x.size
    x_shape = x.shape
    # 初始化导数张量，形状为输入形状拼接输入形状
    df = np.zeros((x_size, x_size), dtype=x.dtype)
    # 填充展平后的对角矩阵
    diag_idx = np.diag_indices(x_size)
    df[diag_idx] = 2 * x.flatten()
    # 重构为对应维度的导数张量
    df = df.reshape(x_shape * 2)
    return f, df

直接用原有的三个场景测试，输出和分场景实现完全一致：

# 标量输入测试
f1, df1 = foo(3.0)
print(f1, df1)  # 输出: 9.0 6.0，匹配foo_scalar结果

# 向量输入测试
x_vec = np.array([1,2,3])
f2, df2 = foo(x_vec)
print(f2)  # 输出: [1 4 9]
print(df2)
# 输出:
# [[2. 0. 0.]
#  [0. 4. 0.]
#  [0. 0. 6.]]  匹配foo_vector结果

# 矩阵输入测试
x_mat = np.array([[1,2],[3,4]])
f3, df3 = foo(x_mat)
print(f3)
# 输出:
# [[ 1  4]
#  [ 9 16]]
print(df3.shape)  # 输出: (2, 2, 2, 2)，匹配foo_matrix的输出形状和值

该实现天然支持更高阶张量输入，例如输入形状为(2,3,4)的三维张量，导数会自动生成形状为(2,3,4,2,3,4)的六阶张量，不需要修改任何代码。

如果处理大尺寸张量，提前生成全0数组会占用较多不必要的内存，可以通过广播直接构造对角张量，减少内存开销：

def foo_fast(x):
    x = np.asarray(x)
    f = x ** 2
    x_shape = x.shape
    ndim = x.ndim
    x_size = x.size
    # 构造单位对角张量
    eye_tensor = np.eye(x_size, dtype=x.dtype).reshape(x_shape * 2)
    # 广播填充导数值
    df = eye_tensor * np.expand_dims(x, axis=tuple(range(ndim, ndim*2)))
    return f, df

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