取消广播 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，returns 一个具有最小形状的数组，如果需要可以将其广播回完整大小。例如：

In [4]: z = unbroadcast(y)

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

然后我可以运行 z 上的第三方功能，然后将结果广播回y.shape。

有没有一种方法可以实现依赖于 Numpy 的 public API 的 unbroadcast？如果没有，是否有任何技巧可以产生预期的结果？

Answer 1

我有一个可能的解决方案，所以 post 在这里（但是，如果有人有更好的解决方案，也请随时回复！）。一种解决方案是检查数组的 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)

Answer 2

这可能相当于你自己的解决方案，只是内置了一点。它在 numpy.lib.stride_tricks:

中使用 as_strided

import numpy as np
from numpy.lib.stride_tricks import as_strided

x = np.arange(16).reshape(2,1,8,1)  # shape (2,1,8,1)
y = np.broadcast_to(x,(2,3,8,5))    # shape (2,3,8,5) broadcast

def unbroadcast(arr):
    #determine unbroadcast shape
    newshape = np.where(np.array(arr.strides) == 0,1,arr.shape) # [2,1,8,1], thanks to @Divakar
    return as_strided(arr,shape=newshape)    # strides are automatically set here

z = unbroadcast(x)
np.all(z==x)  # is True

请注意，在我最初的回答中我没有定义函数，生成的 z 数组具有 (64,0,8,0) 作为 strides，而输入具有 (64,64,8,8) .在当前版本中，返回的 z 数组与 x 具有相同的步幅，我想传递和返回数组会强制创建一个副本。无论如何，我们总是可以在 as_strided 中手动设置步幅以在所有情况下获得相同的数组，但这在上述设置中似乎没有必要。

取消广播 Numpy 数组

Un-broadcasting Numpy arrays

python

arrays

numpy

array-broadcasting