NumPy无for/while重写循环函数实现广播运算的方法

正确实现方案 #

1. bcast 函数 #

原始逻辑中x1维度为(I, J, K)，x2维度为(J, K)，x2缺少x1的第一个I维度，仅需要给x2增加前置广播维度即可完成对齐：

import numpy as np
def bcast(x):
    x1, x2 = x
    return (x1 + 4) * (4 * x2[np.newaxis, :, :] - 4)

x2[np.newaxis, :, :]会将x2维度扩展为(1, J, K)，和维度为(I,J,K)的x1计算时，广播机制会自动把第一个维度复制I次，和原始三重循环逻辑完全等价。

2. bcast_ax 函数 #

原始逻辑中x1维度为(I, K)，x2维度为(J, K)，输出维度为(I, J, K)，分别给两个数组补充缺失的广播维度即可：

def bcast_ax(x):
    x1, x2 = x
    return (4 + x1[:, np.newaxis, :]) * (4 * x2[np.newaxis, :, :] - 4)

• x1[:, np.newaxis, :]将x1扩展为维度(I, 1, K)
• x2[np.newaxis, :, :]将x2扩展为维度(1, J, K)
  两者计算时自动广播为(I,J,K)的输出维度，和原始循环逻辑完全匹配。

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