如何在Python中高效实现兼容一维与多维数组的方法？

一、PyTorch Linear层的底层实现逻辑 #

PyTorch的nn.Linear完全不会采用逐元素循环或递归的实现方式，它依赖底层优化的矩阵乘法+形状变换来处理多维张量：

• 当输入是多维张量（比如形状为(N, D1, D2, ..., Din)）时，Linear层会自动将前n-1个维度视为批量维度，仅对最后一个维度（大小匹配in_features）执行线性变换。
• 具体流程是：先把输入张量的前n-1个维度展平合并成一个大的批量维度，得到形状为(N*D1*D2*..., Din)的二维张量；然后和Linear层的权重矩阵（形状(out_features, in_features)）做矩阵乘法；最后再把结果的形状恢复为(N, D1, D2, ..., Dout)。

举个例子：输入形状(2, 3, 4)，in_features=4，out_features=5

1. 输入展平为(6, 4)（2×3=6）
2. 计算6×4矩阵与5×4权重的转置相乘，得到6×5矩阵
3. 将结果恢复为(2, 3, 5)

这种方式完全利用了BLAS/LAPACK等底层库的矩阵乘法优化，效率远高于循环或递归实现。

二、手动实现高效多维线性变换的方案 #

如果要自己实现类似逻辑，直接用张量形状变换+矩阵乘法即可，这是最高效的方式，无需循环或递归：

代码示例 #

import torch

def multi_dim_linear(input_tensor, weight, bias=None):
    # 保存输入的原始形状
    original_shape = input_tensor.shape
    # 展平前n-1个维度（从第0维到倒数第2维）
    flattened_input = input_tensor.flatten(0, -2)
    # 执行核心线性运算
    output = flattened_input @ weight.T
    if bias is not None:
        output += bias
    # 恢复原始的批量维度结构，替换最后一维为输出特征数
    return output.view(*original_shape[:-1], weight.shape[0])

# 验证与PyTorch原生Linear的一致性
input_tensor = torch.randn(2, 3, 4)
weight = torch.randn(5, 4)
bias = torch.randn(5)

custom_output = multi_dim_linear(input_tensor, weight, bias)
torch_linear = torch.nn.Linear(4, 5)
torch_linear.weight.data = weight
torch_linear.bias.data = bias
torch_output = torch_linear(input_tensor)

print(torch.allclose(custom_output, torch_output))  # 输出 True

三、为什么要避免循环/递归？ #

• 效率极差：Python层面的循环会绕开PyTorch的C++底层优化和GPU加速，处理大张量时速度差距可达数十倍甚至上百倍。
• 代码冗余易出错：递归需要处理各种维度边界情况，代码繁琐且容易出现形状不匹配的bug。
• 完全没必要：PyTorch的张量操作天然支持形状变换和广播机制，用简洁的内置操作就能实现相同逻辑。

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