def bs_pp(*shapes):
ml = max(shapes, key=len)
out = list(ml)
for l in shapes:
if l is ml:
continue
for i, x in enumerate(l, -len(l)):
if x != 1 and x != out[i]:
if out[i] != 1:
raise ValueError
out[i] = x
return (*out,)
def bs_mq1(*shapes):
max_rank = max([len(shape) for shape in shapes])
shapes = [[1] * (max_rank - len(shape)) + shape for shape in shapes]
final_shape = [1] * max_rank
for shape in shapes:
for dim, size in enumerate(shape):
if size != 1:
final_size = final_shape[dim]
if final_size == 1:
final_shape[dim] = size
elif final_size != size:
raise ValueError("Cannot broadcast these shapes")
return (*final_shape,)
import numpy as np
def bs_mq2(*shapes):
max_rank = max([len(shape) for shape in shapes])
shapes = np.array([[1] * (max_rank - len(shape)) + shape
for shape in shapes])
shapes[shapes==1] = -1
final_shape = shapes.max(axis=0)
final_shape[final_shape==-1] = 1
return (*final_shape,)
def bs_hp_ww(*shapes):
return np.broadcast(*[np.empty(shape + [0,], int) for shape in shapes]).shape[:-1]
L = [6], [4, 2, 3, 1], [2, 1, 1]
from timeit import timeit
print('pp: ', timeit(lambda: bs_pp(*L), number=10_000)/10)
print('mq 1: ', timeit(lambda: bs_mq1(*L), number=10_000)/10)
print('mq 2: ', timeit(lambda: bs_mq2(*L), number=10_000)/10)
print('hpaulj/ww:', timeit(lambda: bs_hp_ww(*L), number=10_000)/10)
assert bs_pp(*L) == bs_mq1(*L) and bs_pp(*L) == bs_mq2(*L) and bs_pp(*L) == bs_hp_ww(*L)
In [1]: import numpy
In [2]: numpy.broadcast_shapes((6,), (4, 2, 3, 1), (2, 1, 1))
Out[2]: (4, 2, 3, 6)
对于以前的版本,您可以向每个目标形状广播单个0维数组,然后相互广播所有结果:
def broadcast_shapes(*shapes):
base = numpy.array(0)
broadcast1 = [numpy.broadcast_to(base, shape) for shape in shapes]
return numpy.broadcast(*broadcast1).shape
In [120]: shapes = [6], [4, 2, 3, 1], [2, 1, 1]
In [121]: arrs = np.broadcast_arrays(*[np.empty(shape,int) for shape in shapes])
...:
In [122]: [a.shape for a in arrs]
Out[122]: [(4, 2, 3, 6), (4, 2, 3, 6), (4, 2, 3, 6)]
In [124]: np.lib.stride_tricks._broadcast_shape(*[np.empty(shape,int) for shape
...: in shapes])
Out[124]: (4, 2, 3, 6)
In [131]: np.broadcast(*[np.empty(shape,int) for shape in shapes]).shape
Out[131]: (4, 2, 3, 6)
5条答案
按热度按时间sdnqo3pr1#
下面是另一个直接实现,它碰巧在这个示例中击败了其他实现。最受尊敬的是@hpaulj和@Warren Weckesser的hack,它几乎同样快,也更简洁:
样例运行:
9njqaruj2#
这里有一个简单的实现,以防有人需要它(它可能有助于理解广播)。不过,我更喜欢使用NumPy函数。
编辑
我将这个函数与其他几个答案进行了计时,结果证明它是最快的(编辑,Paul Panzer写了一个更快的函数,参见他的答案,我将其添加到下面的列表中):
yzuktlbb3#
从NumPy 1.20开始,有一个
numpy.broadcast_shapes
函数可以实现您想要的功能。(在文档中,它接受元组而不是列表,所以为了安全起见,您可能应该向它传递元组,但实际上它接受列表。)对于以前的版本,您可以向每个目标形状广播单个0维数组,然后相互广播所有结果:
这避免了为大型形状分配大量内存。不过,需要创建数组感觉有点愚蠢。
hgqdbh6s4#
第二次更快,4.79微秒对42.4微秒。第三次稍微快一点。
正如我最初评论的那样,我从
broadcast_arrays
开始,并查看了代码。然后是_broadcast_shape
,然后是np.broadcast
。9bfwbjaz5#
假设这些形状真的可以广播,那么这个方法是可行的:
如果假设没有空维度,那么
-1
攻击是不必要的: