Numpy中如何实现数组的正交广播（无关维度）

首先，我理解你想要的是不用手动添加固定数量的np.newaxis，而是根据数组A的维度自动调整数组B的形状，让二者可以执行你示例中的广播运算。这里分两种常见的“正交”广播场景来解决：

场景1：保持B的原始维度与A对齐，在B末尾添加与A维度数相同的新轴 #

这就是你示例中展示的用法：给B添加若干新轴到末尾，让A可以广播到这些新轴上。通用实现不需要手动数维度，我们可以通过动态生成新轴的方式来完成：

import numpy as np

# 示例数组
A1 = np.zeros((2,2))
A2 = np.zeros((2,2,2))
B = np.ones((2,2))

# 通用处理函数：给B添加与A维度数相同的新轴到末尾
def expand_b_for_broadcast(A, B):
    # 生成需要添加的新轴位置：全部在末尾
    new_axes = [-1] * A.ndim
    return np.expand_dims(B, axis=new_axes)

# 测试
B_expanded_A1 = expand_b_for_broadcast(A1, B)
C1 = A1 + B_expanded_A1
print(C1.shape)  # 输出: (2, 2, 1, 1)

B_expanded_A2 = expand_b_for_broadcast(A2, B)
C2 = A2 + B_expanded_A2
print(C2.shape)  # 输出: (2, 2, 2, 2, 2)

这里np.expand_dims(B, axis=[-1]*A.ndim)会自动在B的末尾添加A.ndim个新轴，完美适配任意维度的A。

场景2：真正的“正交”广播（A与B的维度完全独立） #

如果你想要的是A的每个元素和B的每个元素进行运算（结果形状为A.shape + B.shape），也就是两个数组的维度完全不重叠，那么需要同时扩展A和B的维度：

import numpy as np

def orthogonal_broadcast(A, B):
    # 给A添加B.ndim个新轴到末尾
    A_expanded = np.expand_dims(A, axis=[-1]*B.ndim)
    # 给B添加A.ndim个新轴到开头
    B_expanded = np.expand_dims(B, axis=[0]*A.ndim)
    return A_expanded + B_expanded

# 测试
C1_ortho = orthogonal_broadcast(A1, B)
print(C1_ortho.shape)  # 输出: (2, 2, 2, 2)

C2_ortho = orthogonal_broadcast(A2, B)
print(C2_ortho.shape)  # 输出: (2, 2, 2, 2, 2)

这种方式下，A的维度和B的维度完全正交，运算时每个A元素会和所有B元素进行配对。

原理说明 #

NumPy的广播规则是从右往左匹配维度：

• 如果两个数组维度数不同，维度少的数组会在左侧自动补1，直到维度数一致
• 对于每个维度，只要其中一个数组的维度是1，就会被广播到另一个数组的维度大小

我们的通用方法就是利用这一规则，通过动态添加1维度来实现任意数组间的适配广播。

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