python - 用 x 计算 x**k，k 是任意维度的数组

Question

我有两个 numpy 数组:一个数组 x 形状为 (n, a0, a1, ...) 和一个数组 k形状为 (n, b0, b1, ...)。我想计算指数数组，使输出具有维度 (a0, a1, ..., b0, b1, ...) 和

out[i0, i1, ..., j0, j1, ...] == prod(x[:, i0, i1, ...] ** k[:, j0, j1, ...])

如果只有一个a_i和一个b_j，广播通过

import numpy

x = numpy.random.rand(2, 31)
k = numpy.random.randint(1, 10, size=(2, 101))

out = numpy.prod(x[..., None]**k[:, None], axis=0)

如果 x 多了几个维度，则必须添加更多的 None:

x = numpy.random.rand(2, 31, 32, 33)
k = numpy.random.randint(1, 10, size=(2, 101))

out = numpy.prod(x[..., None]**k[:, None, None, None], axis=0)

如果x多了几个维度，就需要在其他地方添加更多的None:

x = numpy.random.rand(2, 31)
k = numpy.random.randint(1, 10, size=(2, 51, 51))

out = numpy.prod(x[..., None, None]**k[:, None], axis=0)

如何使 out 的计算与输入数组的维数相关？

Answer 1

下面是使用两个数组的 reshape ，以便它们可以相互广播，然后执行这些操作并沿第一个轴减少prod -

k0_shp = [k.shape[0]] + [1]*(x.ndim-1) + list(k.shape[1:])
x0_shp = list(x.shape) + [1]*(k.ndim-1)
out = (x.reshape(x0_shp) ** k.reshape(k0_shp)).prod(0)

这是将两个输入 reshape 为 3D 的另一种方法，允许每个输入一个单例暗淡，这样它们就可以相互广播，执行 prod 缩减以获得 二维 数组，然后 reshape 回多维度数组 -

s = x.shape[1:] + k.shape[1:] # output shape
out = (x.reshape(x.shape[0],-1,1)**k.reshape(k.shape[0],1,-1)).prod(0).reshape(s)

必须注意， reshape 只是创建输入数组的 View ，因此在内存方面和性能方面实际上都是免费的。