让a 是一个numpy 形状数组(n,m,k) 和a_msk 是一个形状数组(n,m) 包含从a 到掩码的元素乘法.
据我所知，我必须在a_msk 中创建一个新轴，以使其与a 兼容以进行乘法运算。
b = a * a_msk[:,:,np.newaxis]
不幸的是，鉴于数组的大小，我的 Google Colab 运行时在此操作中内存不足。
我的问题是我是否可以在不为掩码数组创建新轴的情况下实现相同的目标。
【问题讨论】：
a 和另一个相同大小的数组。广播不会增加a_msk 的大小，但您需要空间来存储b。
标签： python numpy numpy-ndarray
正如@hpaulj 评论的那样，添加一个轴以使两个数组“兼容”以进行广播是进行乘法的最直接方法。
或者，您可以将数组 a 的最后一个轴移动到前面，这也将使两个数组兼容（我想知道这是否会解决您的内存问题）：
a = np.moveaxis(a, -1, 0)
然后你可以简单地相乘：
b = a * a_msk
但是，要获得结果，您必须将轴移回：
b = np.moveaxis(b, 0, -1)
示例：两种解决方案都返回相同的答案：
import numpy as np
a = np.arange(24).reshape(2, 3, 4)
a_msk = np.arange(6).reshape(2, 3)
print(f'newaxis solution:
{a * a_msk[..., np.newaxis]}')
print()
print(f'moveaxis solution:
{np.moveaxis((np.moveaxis(a, -1, 0) * a_msk), 0, -1)}')
【讨论】：