python - 取消广播 Numpy 数组

Question

在大型代码库中，我使用 np.broadcast_to 来广播数组(这里仅使用简单示例):

In [1]: x = np.array([1,2,3])

In [2]: y = np.broadcast_to(x, (2,1,3))

In [3]: y.shape
Out[3]: (2, 1, 3)

在代码的其他地方，我使用了第三方函数，这些函数可以在 Numpy 数组上以矢量化方式运行，但不是 ufunc。这些函数不理解广播，这意味着在像 y 这样的数组上调用这样的函数是低效的。 Numpy 的 vectorize 等解决方案也不好，因为虽然它们了解广播，但它们在数组元素上引入了一个 for 循环，这样效率很低。

理想情况下，我希望能够拥有一个函数，我们可以调用它，例如unbroadcast，它返回一个具有最小形状的数组，如果需要可以将其广播回完整大小。所以例如:

In [4]: z = unbroadcast(y)

In [5]: z.shape
Out[5]: (1, 1, 3)

然后我可以在 z 上运行第三方函数，然后将结果广播回 y.shape。

有没有一种方法可以实现依赖于 Numpy 公共(public) API 的 unbroadcast？如果没有，是否有任何黑客可以产生预期的结果？

Answer 1

我有一个可能的解决方案，所以会在这里发布(但是如果有人有更好的解决方案，也请随时回复!)。一种解决方案是检查数组的 strides 参数，它在广播维度上将为 0:

def unbroadcast(array):
    slices = []
    for i in range(array.ndim):
        if array.strides[i] == 0:
            slices.append(slice(0, 1))
        else:
            slices.append(slice(None))
    return array[slices]

这给出:

In [14]: unbroadcast(y).shape
Out[14]: (1, 1, 3)